From e0eb4570c664328fa6d8f20366087ea104df7ce4 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 30 Jul 2026 17:07:46 +0200 Subject: [PATCH 01/20] feat(moe): add heterogeneous NVMe expert streaming tier --- server/CMakeLists.txt | 69 + server/docs/DS4.md | 13 + server/docs/ENVIRONMENT.md | 2 + server/docs/KIMI_K3_HETERO.md | 126 ++ server/docs/MOE_NVME_STREAMING.md | 144 ++ server/docs/moe_hybrid.md | 9 + server/src/common/gpu_runtime_compat.h | 2 + server/src/common/moe_hybrid_storage.cpp | 17 +- server/src/common/moe_hybrid_storage.h | 11 +- server/src/common/moe_hybrid_stream.cpp | 1120 ++++++++++---- server/src/common/moe_hybrid_stream.h | 129 +- server/src/common/moe_hybrid_types.h | 13 + server/src/common/moe_nvme_scheduler.cpp | 1636 ++++++++++++++++++++ server/src/common/moe_nvme_scheduler.h | 231 +++ server/src/deepseek4/deepseek4_backend.cpp | 75 +- server/src/deepseek4/deepseek4_graph.cpp | 2 +- server/src/deepseek4/deepseek4_loader.cpp | 11 +- server/src/laguna/laguna_backend.cpp | 22 +- server/src/qwen35moe/qwen35moe_backend.cpp | 20 +- server/test/bench_kimi_k3_hetero.cpp | 386 +++++ server/test/bench_moe_nvme_io.cpp | 241 +++ server/test/bench_moe_nvme_pipeline.cpp | 223 +++ server/test/test_moe_nvme_scheduler.cpp | 364 +++++ 23 files changed, 4498 insertions(+), 368 deletions(-) create mode 100644 server/docs/KIMI_K3_HETERO.md create mode 100644 server/docs/MOE_NVME_STREAMING.md create mode 100644 server/src/common/moe_nvme_scheduler.cpp create mode 100644 server/src/common/moe_nvme_scheduler.h create mode 100644 server/test/bench_kimi_k3_hetero.cpp create mode 100644 server/test/bench_moe_nvme_io.cpp create mode 100644 server/test/bench_moe_nvme_pipeline.cpp create mode 100644 server/test/test_moe_nvme_scheduler.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index ddef3cf44..6e1c0950d 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -306,6 +306,7 @@ add_library(dflash_common STATIC src/common/moe_hybrid_storage.cpp src/common/spark_corpus.cpp src/common/moe_hybrid_ffn_eval.cpp + src/common/moe_nvme_scheduler.cpp src/common/moe_hybrid_stream.cpp src/common/moe_expert_compute.cpp src/common/moe_expert_compute_cpu.cpp @@ -620,6 +621,10 @@ target_link_libraries(dflash_common ggml-base nlohmann_json::nlohmann_json ) +# The NVMe scheduler uses a bounded worker pool on platforms where io_uring is +# unavailable (and as its portable fallback). +find_package(Threads REQUIRED) +target_link_libraries(dflash_common PUBLIC Threads::Threads) # OpenMP for parallel MoE expert compute kernel (saturate memory bandwidth). find_package(OpenMP) if(OpenMP_CXX_FOUND) @@ -669,6 +674,70 @@ if(DFLASH27B_TESTS) add_test(NAME platform_compat COMMAND test_platform_compat) endif() + # SSD scheduling is deliberately independent of ggml and the GPU runtime. + # Keep its correctness tests lightweight so queueing/cache invariants are + # exercised even on CPU-only CI workers. + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_moe_nvme_scheduler.cpp") + add_executable(test_moe_nvme_scheduler + test/test_unit_main.cpp + test/test_moe_nvme_scheduler.cpp + src/common/moe_nvme_scheduler.cpp) + target_include_directories(test_moe_nvme_scheduler PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/src/common + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + target_link_libraries(test_moe_nvme_scheduler PRIVATE Threads::Threads) + list(APPEND _raw_unit_test_targets test_moe_nvme_scheduler) + endif() + + # Read-only microbenchmark for tuning the exact expert I/O path on the + # deployment SSD. It accepts any sufficiently large file; no model parser + # or generated benchmark data is required. + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_moe_nvme_io.cpp") + add_executable(bench_moe_nvme_io + test/bench_moe_nvme_io.cpp + src/common/moe_nvme_scheduler.cpp) + target_include_directories(bench_moe_nvme_io PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/src/common + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + target_link_libraries(bench_moe_nvme_io PRIVATE Threads::Threads) + endif() + + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_moe_nvme_pipeline.cpp") + add_executable(bench_moe_nvme_pipeline test/bench_moe_nvme_pipeline.cpp) + target_include_directories(bench_moe_nvme_pipeline PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/src/common + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + target_link_libraries(bench_moe_nvme_pipeline PRIVATE + dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(bench_moe_nvme_pipeline PRIVATE CUDA::cudart) + else() + target_link_libraries(bench_moe_nvme_pipeline PRIVATE hip::host) + endif() + endif() + + # Kimi-K3 routed-core qualification. This replays the exact latent-expert + # geometry (IQ1_S, 3584 -> 3072 -> 3584, SiTU) through the common NVMe + # stream engine. It is deliberately independent of a Kimi model loader so + # the hardware path can be qualified before downloading the 594 GB model. + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_kimi_k3_hetero.cpp") + add_executable(bench_kimi_k3_hetero test/bench_kimi_k3_hetero.cpp) + target_include_directories(bench_kimi_k3_hetero PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/src/common + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + target_link_libraries(bench_kimi_k3_hetero PRIVATE + dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(bench_kimi_k3_hetero PRIVATE CUDA::cudart) + else() + target_link_libraries(bench_kimi_k3_hetero PRIVATE hip::host) + endif() + endif() + if(DFLASH27B_GPU_BACKEND STREQUAL "hip" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_cuda_pool_shutdown.cpp") add_executable(test_cuda_pool_shutdown diff --git a/server/docs/DS4.md b/server/docs/DS4.md index e0fc2dba7..388f89445 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -175,6 +175,17 @@ If DeepSeek4 is started without an explicit target layer split, `DeepSeek4LayerS The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. +### NVMe cold-capacity tier + +When the cold expert stack cannot fit on Strix, the inference engine turns the +safe remaining Strix memory into an adaptive warm-expert cache and streams only +exact routed misses from NVMe. +`DFLASH_MOE_NVME_COLD_TIER=auto` selects this only when required by measured +free memory; `on` forces qualification and `off` requires resident experts. +The R9700 continues to own dense layers and hot experts. See +[`MOE_NVME_STREAMING.md`](MOE_NVME_STREAMING.md) for the data path, tuning, +and benchmark methodology. + ## Environment Variables | Variable | Purpose | @@ -186,6 +197,8 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. | `DFLASH_DS4_MOE_TP` | Enable routed-expert partitioning. | | `DFLASH_DS4_MOE_TP_INPROC` | Use two local HIP backends instead of an expert IPC worker. | | `DFLASH_DS4_MOE_TP_GPU` | HIP device that owns the cold expert stack. | +| `DFLASH_MOE_NVME_COLD_TIER` | `auto`, `on`, or `off` for the Strix-backed SSD cold-capacity tier. | +| `DFLASH_MOE_NVME_DEVICE_CACHE_MB` | Optional explicit Strix adaptive expert-cache budget; auto mode otherwise uses safe free memory. | | `DFLASH_EXPERT_BUDGET_MB` | Main-GPU memory budget for hot experts. | | `DFLASH_DS4_HOTNESS_CSV` | Optional per-layer routing profile for hot placement. | | `GGML_CUDA_BATCH_PEER_COPIES` | Batch ordered peer copies behind one dependency. | diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index a85e1fa65..e57ddb7e1 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -28,6 +28,8 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_MMID_GROUPED_DEVICE` | -1 | Optional zero-based device restriction; unset/-1 applies to every eligible device. | | `DFLASH_DS4_MOE_TP` / `DFLASH_DS4_MOE_TP_INPROC` | unset | BURN-IN: enable DeepSeek4 route-owner expert parallelism in one process. | | `DFLASH_DS4_MOE_TP_GPU` | auto | HIP device for the cold DeepSeek4 expert owner. | +| `DFLASH_MOE_NVME_COLD_TIER` | auto | BURN-IN: DeepSeek4 cold-capacity policy (`auto`, `on`, `off`). Auto streams only when the cold stack does not fit on Strix with reserve. | +| `DFLASH_MOE_NVME_*` | tuned defaults | BURN-IN: bounded MoE SSD scheduler/backend controls; see `MOE_NVME_STREAMING.md`. | | `GGML_CUDA_BATCH_PEER_COPIES` | unset | BURN-IN: publish ordered HIP peer copies with one cross-device dependency per source/destination pair. | | `DFLASH_MOE_PREFILL_PERSISTENT_OWNER_ALLOC` | 1 for qualified long heterogeneous prefill | KILL SWITCH: =0 restores per-layer route/owner scratch allocation. | | `DFLASH_MOE_TP_*` / `DFLASH_MOE_HYBRID_PREFILL_EAGER` | unset | BURN-IN: model-neutral names for common heterogeneous-MoE scheduling and kernel policy. Existing `DFLASH_DS4_*` names remain compatibility aliases. | diff --git a/server/docs/KIMI_K3_HETERO.md b/server/docs/KIMI_K3_HETERO.md new file mode 100644 index 000000000..ec4e45500 --- /dev/null +++ b/server/docs/KIMI_K3_HETERO.md @@ -0,0 +1,126 @@ +# Kimi K3 on heterogeneous Lucebox + +This note separates measured facts from estimates. The qualification target is +Unsloth `Kimi-K3-UD-IQ1_S`, currently the smallest published Kimi K3 GGUF: 14 +files and 594 GB. + +## What is implemented + +The common MoE data plane now has the Kimi-specific capabilities that were +missing without making the scheduler Kimi-specific: + +- A tensor region carries a model-shard index. `MoeNvmeScheduler` can register + and read any number of GGUF shard file descriptors with `io_uring`, including + an expert whose gate, up, and down tensors live in different files. +- The streamed expert dimension can differ from the model hidden dimension. + This represents Kimi's `7168 -> 3584` routed projection without pretending + its experts consume full-width hidden states. +- The common streamed graph supports SiTU as well as SwiGLU and DS4's clamped + SwiGLU. +- `bench_kimi_k3_hetero` runs Kimi's exact IQ1_S routed-expert geometry through + a reusable graph: 896 experts, top-16, 92 MoE layers, + `3584 -> 3072 -> 3584`, and SiTU (`beta=4`, `linear_beta=25`). It overlaps + SSD/H2D for expert N+1 with compute for expert N. + +Six scheduler tests pass on the AMD Lucebox, including mmap and real-file +multi-shard reads. Existing single-file DS4 descriptors remain source index +zero and need no model-specific change. + +This is not yet a complete Kimi K3 backend. KDA/MLA, Attention Residuals, the +vision encoder, the latent projections around the routed core, tokenizer, and +sampling still need a model adapter. The current upstream llama.cpp Kimi K3 +text-model implementation is also not merged, so it should be treated as a +reference implementation rather than a stable dependency. + +## Exact routed-weight demand + +IQ1_S stores 50 bytes for every 256 values. All three matrices of one Kimi K3 +routed expert therefore occupy: + +```text +gate = row_size(IQ1_S, 3584) * 3072 = 2,150,400 bytes +up = row_size(IQ1_S, 3584) * 3072 = 2,150,400 bytes +down = row_size(IQ1_S, 3072) * 3584 = 2,150,400 bytes +expert = 6,451,200 bytes +``` + +That produces the following deployment constants: + +| Quantity | Exact value | +|---|---:| +| One routed expert | 6.152344 MiB | +| Routed expert calls per decoded token | 92 x 16 = 1,472 | +| Fully cold bytes per token | 8.843994 GiB | +| Complete routed-expert pool | 495.263672 GiB | +| Approximate non-routed part of the 594 GB GGUF | 57.94 GiB | + +The routed pool alone is much larger than Lucebox memory. SSD streaming is +therefore required for this published quant; ordinary CPU/GPU offload cannot +make it resident. + +## Measured result, 2026-07-30 + +The benchmark ran read-only on the R9700 + Strix Halo Lucebox and its P310 +NVMe. It used the exact Kimi byte ranges and IQ1_S+SiTU graph. An existing large +local model supplied the bytes so downloading Kimi's 594 GB was not required; +the byte values do not affect transfer volume or kernel shape. + +| Owner of streamed experts | Scenario | Pipeline | Routed-core rate | +|---|---|---:|---:| +| Strix Halo | 3 tokens, balanced cold routes, compute on | **3.716 GiB/s** | **0.420 token/s** | +| R9700 | 1 token, balanced cold routes, compute on | 2.091 GiB/s | 0.236 token/s | +| Strix Halo | 2 tokens, 10 GiB cache, unrelated routes | 3.638 GiB/s, 0.82% hits | 0.415 token/s | +| Strix Halo | 2 tokens, 10 GiB cache, identical routes | 3.557 GiB/s, 50% aggregate hits | 0.804 token/s | + +The sustained Strix run moved 26.531982 GiB in 7.139575 seconds, evaluated +4,416 routed experts, and reported zero I/O errors. Adding the exact expert +math did not reduce the cold result materially: compute is hidden behind the +8.84 GiB/token storage path. The R9700 is the wrong cold owner because the +extra discrete-GPU upload path cuts end-to-end throughput. + +The repeated-route result is deliberately a best case, not a prediction. +Kimi K3 was designed for balanced expert use. With unrelated balanced routes, +a cache only helps in proportion to its share of the 495 GiB routed pool. + +## Practical Lucebox placement + +The machine has about 125.08 GiB of system/UMA memory plus 31.86 GiB on the +R9700, or 156.94 GiB of unique physical weight capacity before runtime +reserves. A simple placement is: + +1. R9700: attention/KDA/MLA and other dense matrices that fit its 32 GiB. +2. Strix/system memory: remaining non-routed weights, shared experts, latent + projections, recurrent/KV state, workspace, and the routed-expert cache. +3. NVMe: all routed expert stacks, with actual route misses read directly into + pinned slots and evaluated on Strix. + +After approximately 57.94 GiB of non-routed weights plus OS, workspace, and a +moderate context reserve, roughly 70-85 GiB may remain for routed experts. +Under a uniform balanced-routing assumption this covers about 14-17% of the +routed pool. At the measured 3.716 GiB/s, the storage-only ceiling is then +approximately 0.49-0.50 token/s. Full inference will be lower unless dense +R9700 work overlaps almost completely with Strix expert service. + +So the honest expectation for this quant is **roughly one token every two to +three seconds**, not interactive multi-token-per-second generation. Real +router locality can move that estimate; only a route trace from the real model +can establish it. + +## Next end-to-end milestone + +The next useful step is one narrow Kimi adapter, not another generic cache: + +1. Import the correctness-first text graph from the upstream Kimi K3 work. +2. Populate shard-indexed expert regions directly from the GGUF tensor table. +3. Place dense attention on R9700 and the latent/shared/routed MoE path on + Strix; keep only activation-sized transfers at the boundary. +4. Record real `(layer, expert)` routes on a calibration prompt suite and let + the existing placement planner allocate the measured best cache under the + chosen context budget. +5. Compare end-to-end output and token rate against ordinary llama.cpp + CPU/GPU offload. + +The full 594 GB model cannot currently be staged on the qualification box, +which has substantially less free SSD space. It needs at least about 650 GB of +safe free space for all shards plus conversion/logging headroom; existing user +models should not be deleted implicitly. diff --git a/server/docs/MOE_NVME_STREAMING.md b/server/docs/MOE_NVME_STREAMING.md new file mode 100644 index 000000000..b0e1e8557 --- /dev/null +++ b/server/docs/MOE_NVME_STREAMING.md @@ -0,0 +1,144 @@ +# Routed-MoE NVMe Streaming + +This is an inference-engine feature. It streams existing expert-weight bytes +from SSD without changing their format or numerical representation. + +## Lucebox data path + +The three owners have distinct jobs: + +1. **R9700** owns dense layers and the statically hot routed experts. +2. **Strix Halo** owns an adaptive warm-expert cache, the bounded SSD staging + buffers, and execution of streamed cold experts. Its direct path to system + memory is faster than staging those experts into the discrete GPU on the + qualified machine. +3. **NVMe** is the capacity tier for true cache misses that do not fit in the + Strix safe-memory budget. + +`LayerExpertRegions` is the model-adapter contract. It describes the exact +GGUF byte ranges for each layer's gate/up/down tensors (or fused gate+up). +The scheduler therefore has no model names, tensor names, expert dimensions, +or quantization-format assumptions. Each tensor range can also select a model +shard; single-file models use shard zero. See [KIMI_K3_HETERO.md](KIMI_K3_HETERO.md) +for the 14-shard Kimi K3 qualification. + +## Scheduler + +`MoeNvmeScheduler` provides a bounded asynchronous data plane: + +- Linux `io_uring` with a registered model file and fixed, page-locked host + buffers; a portable `pread` worker-pool and mmap fallback remain available. +- Optional `O_DIRECT`, selected automatically when the model is larger than + 75% of physical RAM, avoids retaining both a model-sized page cache and the + explicit expert cache. +- Exact per-expert reads. Direct-I/O requests use aligned envelopes while GPU + copies contain only the logical tensor payload. +- Demand requests outrank speculation. A fixed demand-slot reserve prevents + prefetches from occupying the entire cache. +- Duplicate in-flight requests are merged and speculative requests can be + upgraded to demand without issuing a second read. +- A small LFRU-style resident cache protects demand-loaded experts. A + move-only lease prevents eviction until the asynchronous host-to-device + copy has completed. +- Each expert becomes available at its final completion event. It does not + wait behind the rest of an `io_uring` batch, so SSD read N+1 overlaps the + upload and execution of expert N. +- Two or more rotating GPU slots separate the expert being computed from the + expert being uploaded. +- Otherwise-unused Strix memory becomes a contiguous model-neutral expert + cache indexed by `(layer, expert)`. Cache hits issue neither SSD reads nor + host-to-device copies. LFRU replacement cannot evict a pending upload or an + expert currently executing. + +The native model router remains authoritative. Prediction may only issue a +bounded prefetch; a wrong prediction cannot change model output. + +## DeepSeek V4 Flash activation + +The existing heterogeneous mode is required. In `auto` mode, SSD streaming is +enabled only when the cold expert stack exceeds current Strix free memory +minus the larger of 2 GiB or 5% of device memory. + +```bash +export DFLASH_DS4_MOE_TP=1 +export DFLASH_DS4_MOE_TP_INPROC=1 +export DFLASH_DS4_MOE_TP_GPU=1 +export DFLASH_EXPERT_BUDGET_MB=11700 +export DFLASH_MOE_NVME_COLD_TIER=auto + +./build-hip-dual/dflash_server /path/to/model.gguf \ + --target-device hip:0 --peer-access +``` + +`DFLASH_MOE_NVME_COLD_TIER=on` forces the capacity tier for qualification; +`off` requires the old resident-cold path. In `auto`, a model that exceeds +Strix receives all currently usable memory (after reserve) as its adaptive +expert-cache budget, and only the remainder spills to SSD. + +## Tuning and diagnostics + +Defaults are intentionally small: eight pinned host slots, four fallback I/O +threads, two reserved demand slots, and two GPU slots. More queue depth did +not improve the qualified P310 drive and consumes extra pinned/system memory. + +| Variable | Default | Meaning | +|---|---:|---| +| `DFLASH_MOE_NVME_BACKEND` | `auto` | `auto`, `uring`, `pread`, or `mmap` | +| `DFLASH_MOE_NVME_DIRECT` | `auto` | `auto`, `on`, or `off` | +| `DFLASH_MOE_NVME_SLOTS` | `8` | Fixed pinned host slots | +| `DFLASH_MOE_NVME_IO_THREADS` | `4` | Portable pread workers | +| `DFLASH_MOE_NVME_DEMAND_RESERVE` | `2` | Slots unavailable to speculation | +| `DFLASH_MOE_NVME_PREFETCH_BATCH` | `2` | Maximum speculative jobs per ring submission | +| `DFLASH_MOE_NVME_DEVICE_SLOTS` | `2` | Minimum rotating GPU expert buffers | +| `DFLASH_MOE_NVME_DEVICE_CACHE_MB` | automatic/`0` | Override adaptive Strix expert-cache memory; `0` leaves only pipeline slots | + +Shutdown telemetry reports logical and physical bytes, measured read service +rate, cache hits, demand wait, de-duplication, dropped speculation, and errors. +The standalone targets `test_moe_nvme_scheduler`, `bench_moe_nvme_io`, and +`bench_moe_nvme_pipeline` test correctness, raw storage, and the complete +SSD-to-GPU path respectively. Benchmarks are read-only. + +## Qualification result (2026-07-30) + +On the Lucebox P310 and a realistic 24 MiB expert working set: + +| Path | Direct-I/O throughput | +|---|---:| +| raw SSD scheduler | 4.202 GiB/s | +| SSD to R9700 pipeline | 2.940 GiB/s | +| SSD to Strix pipeline | **4.275 GiB/s** | + +The Strix path completed 182.39 experts/s with p50 42.279 ms and p95 44.069 +ms per eight-expert group and zero I/O errors. The early-completion pipeline +raised the R9700 result from 1.887 to 2.940 GiB/s. These are storage-pipeline +measurements, not end-to-end token rates. + +The adaptive-cache qualification used a 1.584 GiB Strix cache with 64 experts +of 24 MiB each across four rounds. It produced exactly 64 cold misses followed +by 192 GPU hits (75% hit rate), raised effective service to 723.9 expert +accesses/s, and issued only 1.5 GiB of SSD traffic instead of 6 GiB. A separate +16-slot stress case completed 112 evictions at 4.431 GiB/s active I/O with zero +errors. + +With the real 95.3 GiB DeepSeek V4 Flash ROCmFP2 model, an 8 GiB Strix cache +and every HTTP/disk prefix cache disabled, the first five-token prompt loaded +6.068 GiB of unique cold-expert data in 2.406 seconds. The identical second +prompt issued no additional SSD reads and completed in 0.842 seconds: a 2.86x +warm-request improvement with identical output. Final counters were 731 cache +misses, 1,463 Strix hits, and zero I/O errors. + +## Research lineage and next optimization + +The bounded priority/cache design follows the lessons of MoE-Infinity and +HOBBIT; exact per-prefill loading and frequency/recency admission follow +FlashMoE. SpecPrefetch motivates a future shared next-layer transfer predictor +that leaves the native router untouched. MoE-SpAc motivates compile-time +expert layout and I/O coalescing. Tutti's slack-aware `io_uring` scheduling is +relevant when persistent KV traffic shares the device. + +The Kimi qualification now has a persistent reusable expert graph and proves +that its small expert math can hide under SSD service. The production DS4 path +still constructs a graph per streamed expert. Its next implementation step is +to move that qualification design into the common evaluator, then overlap +hot-owner work with cold-owner streaming. Any learned predictor comes after +that deterministic path is qualified. diff --git a/server/docs/moe_hybrid.md b/server/docs/moe_hybrid.md index 5fc7bc131..632271f50 100644 --- a/server/docs/moe_hybrid.md +++ b/server/docs/moe_hybrid.md @@ -23,6 +23,8 @@ The same mechanism supports GPU+CPU offload on a memory-constrained card and GPU | `moe_hybrid_placement.{h,cpp}` | Hot/cold assignment: greedy budget allocation from stats | | `moe_hybrid_swap_manager.{h,cpp}` | Runtime expert promotion/demotion between requests | | `moe_hybrid_storage.{h,cpp}` | Compact owner-local buffers for split expert tensors | +| `moe_nvme_scheduler.{h,cpp}` | Bounded priority SSD I/O, exact expert cache, and leases | +| `moe_hybrid_stream.{h,cpp}` | Pipelined pinned-host to GPU staging and streamed expert evaluation | | `moe_hybrid_ffn_eval.{h,cpp}` | Concurrent owner execution and partial-result joining | | `moe_expert_compute.{h,cpp}` | Backend-neutral selected-expert compute interface | | `moe_expert_compute_ipc.cpp` | Second-GPU process transport and architecture adapter registry | @@ -36,12 +38,14 @@ Model-agnostic architecture descriptor: ```cpp struct MoeHybridConfig { int n_embd; // hidden dimension + int n_expert_embd; // routed latent dim; 0 means n_embd int n_expert; // total experts per layer int n_expert_used; // top-k selected per token int n_ff_exp; // routed expert intermediate dim int n_ff_shexp; // shared expert intermediate dim (0 = none) int n_layer; // number of MoE layers int first_moe_layer; // first MoE layer index + MoeGatedActivation gated_activation; // SwiGLU or SiTU }; ``` @@ -153,6 +157,11 @@ Two loading paths: 2. **From file** (`build_moe_hybrid_storage_from_file`): Reads expert slices directly from mmap'd GGUF data into the selected owners, avoiding a temporary full expert stack on either GPU. +3. **NVMe capacity tier**: Leaves non-resident cold tensors in the model file, + retains exact per-layer byte descriptors, and loads only routed misses into + a bounded pinned-host pipeline plus an adaptive complementary-GPU cache. See + [`MOE_NVME_STREAMING.md`](MOE_NVME_STREAMING.md). + Both paths produce the same `MoeHybridStorage` containing per-layer split buffers ready for evaluation. ## Model Integration Contract diff --git a/server/src/common/gpu_runtime_compat.h b/server/src/common/gpu_runtime_compat.h index 0cfc9d76b..6263e24aa 100644 --- a/server/src/common/gpu_runtime_compat.h +++ b/server/src/common/gpu_runtime_compat.h @@ -44,6 +44,8 @@ #define cudaDeviceReset hipDeviceReset #define cudaEvent_t hipEvent_t #define cudaEventCreate hipEventCreate +#define cudaEventCreateWithFlags hipEventCreateWithFlags +#define cudaEventDisableTiming hipEventDisableTiming #define cudaEventDestroy hipEventDestroy #define cudaEventElapsedTime hipEventElapsedTime #define cudaEventRecord hipEventRecord diff --git a/server/src/common/moe_hybrid_storage.cpp b/server/src/common/moe_hybrid_storage.cpp index b1ad20f35..65a96448e 100644 --- a/server/src/common/moe_hybrid_storage.cpp +++ b/server/src/common/moe_hybrid_storage.cpp @@ -16,6 +16,7 @@ #if !defined(_WIN32) #include +#include #else #if !defined(NOMINMAX) #define NOMINMAX @@ -212,6 +213,12 @@ MoeHybridStorage::~MoeHybridStorage() { mmap_data = nullptr; mmap_size = 0; } +#if !defined(_WIN32) + if (mmap_fd >= 0) { + ::close(mmap_fd); + mmap_fd = -1; + } +#endif } bool MoeHybridStorage::matches(const MoeHybridConfig & cfg) const { @@ -745,7 +752,8 @@ bool build_moe_hybrid_storage_from_file_with_mmap( MoeHybridStorage & out, std::string * err, int cache_slots, - ggml_backend_t cold_gpu_backend) { + ggml_backend_t cold_gpu_backend, + int mmap_fd) { // First build storage normally (hot GPU + cold CPU buffers). if (!build_moe_hybrid_storage_from_file( @@ -757,6 +765,13 @@ bool build_moe_hybrid_storage_from_file_with_mmap( // Store mmap metadata for streaming prefill. out.mmap_data = mmap_base; out.mmap_size = mmap_total_size; +#if !defined(_WIN32) + // Keep a separate open-file description for true asynchronous reads. A + // failure is non-fatal: mmap workers remain a correct portable fallback. + if (mmap_fd >= 0) out.mmap_fd = ::dup(mmap_fd); +#else + (void) mmap_fd; +#endif // Compute per-layer expert file regions (offsets relative to mmap base). const auto * base = static_cast(mmap_base); diff --git a/server/src/common/moe_hybrid_storage.h b/server/src/common/moe_hybrid_storage.h index 12c8fb83d..e8558fb57 100644 --- a/server/src/common/moe_hybrid_storage.h +++ b/server/src/common/moe_hybrid_storage.h @@ -16,10 +16,12 @@ namespace dflash::common { struct MoeHybridRoutingStats; -// File region for one expert tensor (offset into mmap). +// File region for one expert tensor. source_index is zero for ordinary +// single-file models and selects a shard for split GGUF models. struct ExpertFileRegion { size_t offset = 0; size_t size = 0; + uint32_t source_index = 0; }; // Per-layer file regions for all expert tensors (used by streaming prefill). @@ -211,7 +213,9 @@ struct MoeHybridStorage { // When set, the streaming engine can DMA cold experts directly from here. const void * mmap_data = nullptr; size_t mmap_size = 0; - int mmap_fd = -1; // POSIX fd for madvise; -1 on Windows or if not available + // Owned duplicate of the model fd. The SSD scheduler duplicates it again, + // allowing io_uring/pread/O_DIRECT without coupling its lifetime to mmap. + int mmap_fd = -1; // Per-layer file region metadata for streaming (populated when mmap is active). std::vector layer_regions; @@ -287,6 +291,7 @@ bool build_moe_hybrid_storage_from_file_with_mmap( MoeHybridStorage & out, std::string * err = nullptr, int cache_slots = 0, - ggml_backend_t cold_gpu_backend = nullptr); + ggml_backend_t cold_gpu_backend = nullptr, + int mmap_fd = -1); } // namespace dflash::common diff --git a/server/src/common/moe_hybrid_stream.cpp b/server/src/common/moe_hybrid_stream.cpp index 0b8cd74ff..f3bade16c 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -1,13 +1,18 @@ -// MoE hybrid prefill streaming engine — implementation. - #include "moe_hybrid_stream.h" #include "gpu_runtime_compat.h" +#include "ggml-alloc.h" #include "ggml-backend.h" +#include "ggml-cuda.h" #include #include +#include #include +#include +#include +#include +#include #if !defined(_WIN32) #include @@ -15,46 +20,252 @@ #endif namespace dflash::common { +namespace { -MoeHybridStreamEngine::~MoeHybridStreamEngine() { - destroy(); +bool pinned_allocate(void ** ptr, size_t bytes, void *) { + return cudaMallocHost(ptr, bytes) == cudaSuccess; +} + +void pinned_free(void * ptr, void *) { + if (ptr) (void) cudaFreeHost(ptr); +} + +int env_bounded_int(const char * name, int fallback, int lo, int hi) { + const char * value = std::getenv(name); + if (!value || !value[0]) return fallback; + char * end = nullptr; + const long parsed = std::strtol(value, &end, 10); + if (end == value || *end != '\0' || parsed < lo || parsed > hi) return fallback; + return (int) parsed; +} + +size_t env_mib(const char * name, size_t fallback) { + const char * value = std::getenv(name); + if (!value || !value[0] || value[0] == '-') return fallback; + char * end = nullptr; + const unsigned long long parsed = std::strtoull(value, &end, 10); + constexpr size_t kMiB = 1024 * 1024; + if (end == value || *end != '\0' || + parsed > std::numeric_limits::max() / kMiB) { + return fallback; + } + return (size_t) parsed * kMiB; +} + +size_t align_up(size_t value, size_t alignment) { + if (alignment == 0 || value > std::numeric_limits::max() - (alignment - 1)) { + return 0; + } + return (value + alignment - 1) & ~(alignment - 1); +} + +uint64_t device_key(int layer, int expert) { + return ((uint64_t) (uint32_t) layer << 32) | (uint32_t) expert; } -MoeHybridStreamEngine::MoeHybridStreamEngine(MoeHybridStreamEngine && o) noexcept - : pinned_buf_(o.pinned_buf_), pinned_size_(o.pinned_size_), - gpu_scratch_(o.gpu_scratch_), scratch_size_(o.scratch_size_), - backend_(o.backend_), - scratch_gate_(o.scratch_gate_), scratch_up_(o.scratch_up_), - scratch_down_(o.scratch_down_), - last_gate_bytes_(o.last_gate_bytes_), last_up_bytes_(o.last_up_bytes_), - last_down_bytes_(o.last_down_bytes_) { - o.pinned_buf_ = nullptr; o.pinned_size_ = 0; - o.gpu_scratch_ = nullptr; o.scratch_size_ = 0; - o.backend_ = nullptr; - o.scratch_gate_ = nullptr; o.scratch_up_ = nullptr; o.scratch_down_ = nullptr; - o.last_gate_bytes_ = 0; o.last_up_bytes_ = 0; o.last_down_bytes_ = 0; +int backend_device_index(ggml_backend_t backend) { + if (!backend || !ggml_backend_is_cuda(backend)) return -1; + ggml_backend_dev_t wanted = ggml_backend_get_device(backend); + ggml_backend_reg_t reg = ggml_backend_cuda_reg(); + const int count = ggml_backend_cuda_get_device_count(); + for (int device = 0; device < count; ++device) { + if (ggml_backend_reg_dev_get(reg, (size_t) device) == wanted) return device; + } + return -1; +} + +// HIP/CUDA streams, events, and allocations belong to the current device. +// The heterogeneous engine alternates R9700 and Strix backends on one host +// thread, so relying on whichever backend ran last is a cross-device bug. +class ScopedGpuDevice { +public: + explicit ScopedGpuDevice(int target) : target_(target) { + if (target_ < 0 || cudaGetDevice(&previous_) != cudaSuccess) return; + valid_ = true; + if (previous_ != target_) switched_ = cudaSetDevice(target_) == cudaSuccess; + } + + ~ScopedGpuDevice() { + if (valid_ && switched_) (void) cudaSetDevice(previous_); + } + + bool ready() const { return valid_ && (previous_ == target_ || switched_); } + +private: + int target_ = -1; + int previous_ = -1; + bool valid_ = false; + bool switched_ = false; +}; + +} // namespace + +MoeStreamConfig MoeStreamConfig::from_env() { + MoeStreamConfig config; + config.nvme = MoeNvmeConfig::from_env(config.nvme); + config.device_slots = env_bounded_int( + "DFLASH_MOE_NVME_DEVICE_SLOTS", config.device_slots, 2, 8); + config.device_cache_bytes = env_mib( + "DFLASH_MOE_NVME_DEVICE_CACHE_MB", config.device_cache_bytes); + config.prefill_threshold = env_bounded_int( + "DFLASH_MOE_NVME_PREFILL_THRESHOLD", config.prefill_threshold, 1, 4096); + return config; +} + +struct MoeHybridStreamEngine::Runtime { + struct DeviceSlot { + void * data = nullptr; + cudaEvent_t ready = nullptr; + bool pending = false; + bool valid = false; + bool cache_managed = false; + int compute_users = 0; + MoeExpertKey key{}; + uint64_t frequency = 0; + uint64_t last_touch = 0; + MoeNvmeLease host_lease; + MoeExpertIoLayout layout{}; + }; + + ggml_backend_t backend = nullptr; + int device = -1; + size_t max_expert_bytes = 0; + MoeStreamConfig config{}; + std::unique_ptr io; + cudaStream_t transfer_stream = nullptr; + void * device_pool = nullptr; + size_t device_stride = 0; + size_t device_pool_bytes = 0; + std::vector device_slots; + std::unordered_map device_index; + uint64_t device_clock = 0; + uint64_t device_cache_hits = 0; + uint64_t device_cache_misses = 0; + uint64_t device_cache_evictions = 0; + int active_slot = -1; +}; + +template +bool allocate_device_cache(RuntimeT & runtime, std::string * err) { + runtime.device_stride = align_up(runtime.max_expert_bytes, 256); + if (runtime.device_stride == 0) { + if (err) *err = "SSD device-cache stride overflow"; + return false; + } + + size_t desired_slots = (size_t) std::max(2, runtime.config.device_slots); + if (runtime.config.device_cache_bytes > 0) { + desired_slots = std::max( + desired_slots, runtime.config.device_cache_bytes / runtime.device_stride); + } + constexpr size_t kMaxDeviceSlots = 65536; + desired_slots = std::min(desired_slots, kMaxDeviceSlots); + desired_slots = std::min( + desired_slots, std::numeric_limits::max() / runtime.device_stride); + + // A large contiguous allocation keeps address arithmetic cheap and avoids + // thousands of allocator objects. If the planner's free-memory snapshot + // raced another allocation, converge to a smaller usable cache instead of + // failing model startup. + size_t attempt_slots = desired_slots; + cudaError_t gpu_err = cudaSuccess; + while (attempt_slots >= 2) { + const size_t bytes = attempt_slots * runtime.device_stride; + gpu_err = cudaMalloc(&runtime.device_pool, bytes); + if (gpu_err == cudaSuccess) { + runtime.device_pool_bytes = bytes; + break; + } + if (attempt_slots == 2) break; + attempt_slots = std::max(2, attempt_slots * 3 / 4); + } + if (!runtime.device_pool) { + if (err) { + *err = std::string("failed to allocate SSD GPU expert cache: ") + + cudaGetErrorString(gpu_err); + } + return false; + } + + try { + runtime.device_slots.resize(attempt_slots); + } catch (const std::bad_alloc &) { + (void) cudaFree(runtime.device_pool); + runtime.device_pool = nullptr; + runtime.device_pool_bytes = 0; + if (err) *err = "failed to allocate SSD GPU cache metadata"; + return false; + } + auto * base = static_cast(runtime.device_pool); + for (size_t i = 0; i < runtime.device_slots.size(); ++i) { + runtime.device_slots[i].data = base + i * runtime.device_stride; + } + runtime.config.device_slots = (int) attempt_slots; + return true; } -MoeHybridStreamEngine & MoeHybridStreamEngine::operator=(MoeHybridStreamEngine && o) noexcept { - if (this != &o) { +MoeHybridStreamEngine::MoeHybridStreamEngine() = default; +MoeHybridStreamEngine::~MoeHybridStreamEngine() { destroy(); } +MoeHybridStreamEngine::MoeHybridStreamEngine(MoeHybridStreamEngine &&) noexcept = default; +MoeHybridStreamEngine & MoeHybridStreamEngine::operator=(MoeHybridStreamEngine &&) noexcept = default; + +bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, size_t max_expert_bytes, + std::string * err) { + destroy(); + if (!gpu_backend || max_expert_bytes == 0) { + if (err) *err = "invalid arguments to stream engine init"; + return false; + } + + std::unique_ptr runtime(new (std::nothrow) Runtime); + if (!runtime) { + if (err) *err = "failed to allocate stream runtime"; + return false; + } + runtime->backend = gpu_backend; + runtime->device = backend_device_index(gpu_backend); + ScopedGpuDevice device_scope(runtime->device); + if (!device_scope.ready()) { + if (err) *err = "failed to resolve/select SSD stream GPU"; + return false; + } + runtime->max_expert_bytes = max_expert_bytes; + runtime->config = MoeStreamConfig::from_env(); + runtime->io.reset(new (std::nothrow) MoeNvmeScheduler); + if (!runtime->io) { + if (err) *err = "failed to allocate SSD scheduler"; + return false; + } + if (!runtime->io->init(runtime->config.nvme, max_expert_bytes, + pinned_allocate, pinned_free, nullptr, err)) { + return false; + } + + cudaError_t gpu_err = cudaStreamCreate(&runtime->transfer_stream); + if (gpu_err != cudaSuccess) { + if (err) *err = std::string("failed to create SSD transfer stream: ") + + cudaGetErrorString(gpu_err); + return false; + } + if (!allocate_device_cache(*runtime, err)) { + // Install the partial runtime so destroy() releases every resource. + runtime_ = std::move(runtime); destroy(); - pinned_buf_ = o.pinned_buf_; pinned_size_ = o.pinned_size_; - gpu_scratch_ = o.gpu_scratch_; scratch_size_ = o.scratch_size_; - backend_ = o.backend_; - scratch_gate_ = o.scratch_gate_; scratch_up_ = o.scratch_up_; - scratch_down_ = o.scratch_down_; - last_gate_bytes_ = o.last_gate_bytes_; last_up_bytes_ = o.last_up_bytes_; - last_down_bytes_ = o.last_down_bytes_; - o.pinned_buf_ = nullptr; o.pinned_size_ = 0; - o.gpu_scratch_ = nullptr; o.scratch_size_ = 0; - o.backend_ = nullptr; - o.scratch_gate_ = nullptr; o.scratch_up_ = nullptr; o.scratch_down_ = nullptr; - o.last_gate_bytes_ = 0; o.last_up_bytes_ = 0; o.last_down_bytes_ = 0; - } - return *this; + return false; + } + runtime_ = std::move(runtime); + return true; +} + +bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, size_t max_expert_bytes, + const MoeHybridStorage & storage, + std::string * err) { + return init(gpu_backend, max_expert_bytes, storage, MoeStreamConfig::from_env(), err); } bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, size_t max_expert_bytes, + const MoeHybridStorage & storage, + const MoeStreamConfig & config, std::string * err) { destroy(); if (!gpu_backend || max_expert_bytes == 0) { @@ -62,209 +273,490 @@ bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, size_t max_expert_b return false; } - // Allocate pinned host staging buffer - cudaError_t cuda_err = cudaMallocHost(&pinned_buf_, max_expert_bytes); - if (cuda_err != cudaSuccess) { - if (err) *err = std::string("cudaMallocHost failed: ") + cudaGetErrorString(cuda_err); + std::unique_ptr runtime(new (std::nothrow) Runtime); + if (!runtime) { + if (err) *err = "failed to allocate stream runtime"; + return false; + } + runtime->backend = gpu_backend; + runtime->device = backend_device_index(gpu_backend); + ScopedGpuDevice device_scope(runtime->device); + if (!device_scope.ready()) { + if (err) *err = "failed to resolve/select SSD stream GPU"; + return false; + } + runtime->max_expert_bytes = max_expert_bytes; + runtime->config = config; + runtime->config.device_slots = std::max(2, runtime->config.device_slots); + runtime->io.reset(new (std::nothrow) MoeNvmeScheduler); + if (!runtime->io || + !runtime->io->init(runtime->config.nvme, max_expert_bytes, + pinned_allocate, pinned_free, nullptr, err)) { return false; } - pinned_size_ = max_expert_bytes; - // Allocate GPU scratch buffer - cuda_err = cudaMalloc(&gpu_scratch_, max_expert_bytes); - if (cuda_err != cudaSuccess) { - if (err) *err = std::string("cudaMalloc scratch failed: ") + cudaGetErrorString(cuda_err); - cudaFreeHost(pinned_buf_); - pinned_buf_ = nullptr; - pinned_size_ = 0; + cudaError_t gpu_err = cudaStreamCreate(&runtime->transfer_stream); + if (gpu_err != cudaSuccess) { + if (err) *err = std::string("failed to create SSD transfer stream: ") + + cudaGetErrorString(gpu_err); + return false; + } + if (!allocate_device_cache(*runtime, err)) { + runtime_ = std::move(runtime); + destroy(); + return false; + } + runtime_ = std::move(runtime); + if (!bind_storage(storage, err)) { + destroy(); return false; } - scratch_size_ = max_expert_bytes; - backend_ = gpu_backend; return true; } +bool MoeHybridStreamEngine::bind_storage(const MoeHybridStorage & storage, + std::string * err) { + if (!runtime_ || !runtime_->io || !runtime_->io->is_initialized()) { + if (err) *err = "stream engine is not initialized"; + return false; + } + return runtime_->io->bind_source( + {storage.mmap_data, storage.mmap_size, storage.mmap_fd}, + storage.layer_regions, err); +} + +bool MoeHybridStreamEngine::bind_sources( + const std::vector & sources, + const std::vector & layer_regions, + std::string * err) { + if (!runtime_ || !runtime_->io || !runtime_->io->is_initialized()) { + if (err) *err = "stream engine is not initialized"; + return false; + } + return runtime_->io->bind_sources(sources, layer_regions, err); +} + bool MoeHybridStreamEngine::is_ready() const { - return pinned_buf_ && gpu_scratch_ && backend_; + return runtime_ && runtime_->backend && runtime_->io && + runtime_->io->is_initialized() && runtime_->transfer_stream && + !runtime_->device_slots.empty(); +} + +bool MoeHybridStreamEngine::is_bound() const { + return is_ready() && runtime_->io->is_bound(); } void MoeHybridStreamEngine::destroy() { - if (gpu_scratch_) { - cudaFree(gpu_scratch_); - gpu_scratch_ = nullptr; - } - if (pinned_buf_) { - cudaFreeHost(pinned_buf_); - pinned_buf_ = nullptr; - } - pinned_size_ = 0; - scratch_size_ = 0; - backend_ = nullptr; - scratch_gate_ = nullptr; - scratch_up_ = nullptr; - scratch_down_ = nullptr; - last_gate_bytes_ = 0; - last_up_bytes_ = 0; - last_down_bytes_ = 0; -} - -void MoeHybridStreamEngine::prefetch_cold_experts(const void * mmap_data, size_t mmap_size, - const LayerExpertRegions & regions, - const int32_t * cold_expert_ids, - int n_cold) { + if (!runtime_) return; + const size_t device_cache_slot_count = runtime_->device_slots.size(); + const size_t device_cache_byte_count = runtime_->device_pool_bytes; + ScopedGpuDevice device_scope(runtime_->device); + if (runtime_->transfer_stream) { + (void) cudaStreamSynchronize(runtime_->transfer_stream); + } + for (Runtime::DeviceSlot & slot : runtime_->device_slots) { + slot.host_lease.reset(); + if (slot.ready) (void) cudaEventDestroy(slot.ready); + slot.ready = nullptr; + slot.data = nullptr; + slot.pending = false; + } + runtime_->device_slots.clear(); + runtime_->device_index.clear(); + if (runtime_->device_pool) (void) cudaFree(runtime_->device_pool); + runtime_->device_pool = nullptr; + if (runtime_->transfer_stream) (void) cudaStreamDestroy(runtime_->transfer_stream); + runtime_->transfer_stream = nullptr; + if (runtime_->io) { + const MoeNvmeStats stats = runtime_->io->stats(); + if (stats.requests != 0 || stats.read_ops != 0 || stats.errors != 0) { + const double payload_gib = (double) stats.payload_bytes / + (1024.0 * 1024.0 * 1024.0); + const double physical_gib = (double) stats.physical_bytes / + (1024.0 * 1024.0 * 1024.0); + const double read_seconds = (double) stats.active_io_ns / 1.0e9; + const double read_gib_s = read_seconds > 0.0 + ? physical_gib / read_seconds : 0.0; + const double hit_rate = stats.requests > 0 + ? 100.0 * (double) stats.cache_hits / (double) stats.requests : 0.0; + const double mean_wait_ms = stats.demand_requests > 0 + ? ((double) stats.wait_ns / 1.0e6) / + (double) stats.demand_requests : 0.0; + std::fprintf(stderr, + "[moe-nvme] io=%s requests=%llu reads=%llu " + "payload=%.3f GiB physical=%.3f GiB active-io-rate=%.3f GiB/s " + "cache-hit=%.1f%% mean-demand-wait=%.3f ms " + "dedupe=%llu upgrades=%llu dropped-prefetch=%llu errors=%llu " + "strix-cache=%.1f MiB slots=%zu hits=%llu misses=%llu evictions=%llu\n", + runtime_->io->effective_backend_name(), + (unsigned long long) stats.requests, + (unsigned long long) stats.read_ops, + payload_gib, physical_gib, read_gib_s, hit_rate, mean_wait_ms, + (unsigned long long) stats.inflight_deduplications, + (unsigned long long) stats.demand_upgrades, + (unsigned long long) stats.prefetch_drops, + (unsigned long long) stats.errors, + device_cache_byte_count / 1024.0 / 1024.0, + device_cache_slot_count, + (unsigned long long) runtime_->device_cache_hits, + (unsigned long long) runtime_->device_cache_misses, + (unsigned long long) runtime_->device_cache_evictions); + } + runtime_->io->destroy(); + } + runtime_.reset(); +} + +void MoeHybridStreamEngine::request_experts(int layer, const int32_t * expert_ids, + int count, MoeNvmePriority priority) { + if (!is_bound() || !expert_ids || count <= 0) return; + for (int i = 0; i < count; ++i) { + if (expert_ids[i] < 0) continue; + const uint64_t key = device_key(layer, expert_ids[i]); + const auto cached = runtime_->device_index.find(key); + if (cached != runtime_->device_index.end()) { + const int slot_index = cached->second; + if (slot_index >= 0 && slot_index < (int) runtime_->device_slots.size()) { + const Runtime::DeviceSlot & slot = + runtime_->device_slots[(size_t) slot_index]; + if (slot.valid && slot.key.layer == layer && + slot.key.expert == expert_ids[i]) { + continue; + } + } + runtime_->device_index.erase(cached); + } + (void) runtime_->io->request(layer, expert_ids[i], priority, nullptr); + } +} + +void MoeHybridStreamEngine::prefetch_cold_experts( + const void * mmap_data, size_t mmap_size, + const LayerExpertRegions & regions, + const int32_t * cold_expert_ids, int n_cold) { if (!mmap_data || mmap_size == 0 || !cold_expert_ids || n_cold <= 0) return; #if !defined(_WIN32) - auto do_advise = [&](size_t offset, size_t length) { - if (offset + length > mmap_size) return; - const size_t page_size = (size_t)sysconf(_SC_PAGESIZE); - const size_t aligned_offset = (offset / page_size) * page_size; - const size_t aligned_length = length + (offset - aligned_offset); - ::madvise(const_cast(static_cast(mmap_data)) + aligned_offset, - aligned_length, MADV_WILLNEED); + const size_t page_size = (size_t) std::max(1, ::sysconf(_SC_PAGESIZE)); + auto advise = [&](size_t offset, size_t bytes) { + if (offset > mmap_size || bytes > mmap_size - offset || bytes == 0) return; + const size_t aligned = (offset / page_size) * page_size; + const size_t length = bytes + (offset - aligned); + (void) ::madvise( + const_cast(static_cast(mmap_data)) + aligned, + length, MADV_WILLNEED); }; -#endif - for (int i = 0; i < n_cold; ++i) { - const int32_t eid = cold_expert_ids[i]; -#if !defined(_WIN32) + const int expert = cold_expert_ids[i]; + if (expert < 0) continue; if (regions.fused_gate_up) { - if (regions.gate_up_exps.size > 0) { - do_advise(regions.gate_up_exps.offset + (size_t)eid * regions.expert_bytes_gate_up, - regions.expert_bytes_gate_up); - } + advise(regions.gate_up_exps.offset + (size_t) expert * regions.expert_bytes_gate_up, + regions.expert_bytes_gate_up); } else { - if (regions.gate_exps.size > 0) { - do_advise(regions.gate_exps.offset + (size_t)eid * regions.expert_bytes_gate, - regions.expert_bytes_gate); - } - if (regions.up_exps.size > 0) { - do_advise(regions.up_exps.offset + (size_t)eid * regions.expert_bytes_up, - regions.expert_bytes_up); - } - } - if (regions.down_exps.size > 0) { - do_advise(regions.down_exps.offset + (size_t)eid * regions.expert_bytes_down, - regions.expert_bytes_down); + advise(regions.gate_exps.offset + (size_t) expert * regions.expert_bytes_gate, + regions.expert_bytes_gate); + advise(regions.up_exps.offset + (size_t) expert * regions.expert_bytes_up, + regions.expert_bytes_up); } + advise(regions.down_exps.offset + (size_t) expert * regions.expert_bytes_down, + regions.expert_bytes_down); + } #else - (void)eid; - (void)regions; + (void) regions; #endif - } } -bool MoeHybridStreamEngine::stream_expert_sync(const void * mmap_data, size_t mmap_size, - const LayerExpertRegions & regions, - int expert_id, - ggml_backend_t gpu_backend, - std::string * err) { - if (!is_ready()) { - if (err) *err = "stream engine not initialized"; +bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, + int device_slot, + std::string * err) { + if (!is_bound()) { + if (err) *err = "stream engine has no bound SSD model source"; return false; } - if (!mmap_data || mmap_size == 0) { - if (err) *err = "mmap not available"; + ScopedGpuDevice device_scope(runtime_->device); + if (!device_scope.ready()) { + if (err) *err = "failed to select SSD stream GPU"; return false; } + if (device_slot < 0 || device_slot >= (int) runtime_->device_slots.size()) { + if (err) *err = "SSD device slot is out of range"; + return false; + } + Runtime::DeviceSlot & dst = runtime_->device_slots[(size_t) device_slot]; + if (dst.compute_users != 0) { + if (err) *err = "SSD device slot is still in use by expert compute"; + return false; + } + if (dst.pending) { + const cudaError_t wait_err = cudaEventSynchronize(dst.ready); + if (wait_err != cudaSuccess) { + if (err) *err = std::string("failed waiting for prior SSD upload: ") + + cudaGetErrorString(wait_err); + return false; + } + dst.pending = false; + dst.host_lease.reset(); + } + if (dst.cache_managed && dst.valid) { + runtime_->device_index.erase(device_key(dst.key.layer, dst.key.expert)); + } + dst.valid = false; + dst.cache_managed = false; + dst.key = {}; - const auto * file_base = static_cast(mmap_data); - size_t staging_offset = 0; + if (!dst.ready) { + const cudaError_t event_create_err = + cudaEventCreateWithFlags(&dst.ready, cudaEventDisableTiming); + if (event_create_err != cudaSuccess) { + if (err) *err = std::string("failed to create expert upload event: ") + + cudaGetErrorString(event_create_err); + return false; + } + } - // Validate expert_id against region size - if (expert_id < 0) { - if (err) *err = "expert_id is negative"; + MoeNvmeLease lease; + if (!runtime_->io->acquire(layer, expert_id, lease, err)) return false; + if (lease.layout().payload_bytes > runtime_->max_expert_bytes) { + if (err) *err = "streamed expert exceeds GPU device slot"; return false; } - - // Copy gate (or fused gate_up) from mmap → pinned - if (regions.fused_gate_up) { - const size_t bytes = regions.expert_bytes_gate_up; - const size_t file_off = regions.gate_up_exps.offset + (size_t)expert_id * bytes; - if (file_off + bytes > mmap_size) { - if (err) *err = "gate_up expert out of file bounds"; + for (int i = 0; i < lease.layout().span_count; ++i) { + const MoeExpertIoSpan & span = lease.layout().spans[i]; + cudaError_t gpu_err = cudaMemcpyAsync( + static_cast(dst.data) + span.device_offset, + lease.data() + span.buffer_offset, + span.bytes, cudaMemcpyHostToDevice, runtime_->transfer_stream); + if (gpu_err != cudaSuccess) { + (void) cudaStreamSynchronize(runtime_->transfer_stream); + if (err) *err = std::string("asynchronous expert H2D failed: ") + + cudaGetErrorString(gpu_err); return false; } - std::memcpy(static_cast(pinned_buf_) + staging_offset, - file_base + file_off, bytes); - last_gate_bytes_ = bytes; - last_up_bytes_ = 0; - staging_offset += bytes; - } else { - // gate - { - const size_t bytes = regions.expert_bytes_gate; - const size_t file_off = regions.gate_exps.offset + (size_t)expert_id * bytes; - if (file_off + bytes > mmap_size) { - if (err) *err = "gate expert out of file bounds"; - return false; + } + const cudaError_t event_err = cudaEventRecord(dst.ready, runtime_->transfer_stream); + if (event_err != cudaSuccess) { + (void) cudaStreamSynchronize(runtime_->transfer_stream); + if (err) *err = std::string("failed to record expert upload event: ") + + cudaGetErrorString(event_err); + return false; + } + dst.layout = lease.layout(); + dst.host_lease = std::move(lease); + dst.pending = true; + return true; +} + +bool MoeHybridStreamEngine::stage_expert_cached_async( + int layer, int expert_id, int * device_slot, std::string * err) { + if (!device_slot) { + if (err) *err = "SSD cache stage requires an output slot"; + return false; + } + *device_slot = -1; + if (!is_bound()) { + if (err) *err = "stream engine has no bound SSD model source"; + return false; + } + + const uint64_t key = device_key(layer, expert_id); + auto cached = runtime_->device_index.find(key); + if (cached != runtime_->device_index.end()) { + const int index = cached->second; + if (index >= 0 && index < (int) runtime_->device_slots.size()) { + Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) index]; + if (slot.valid && slot.cache_managed && + slot.key.layer == layer && slot.key.expert == expert_id) { + ++runtime_->device_cache_hits; + ++slot.frequency; + slot.last_touch = ++runtime_->device_clock; + *device_slot = index; + return true; } - std::memcpy(static_cast(pinned_buf_) + staging_offset, - file_base + file_off, bytes); - last_gate_bytes_ = bytes; - staging_offset += bytes; } - // up - { - const size_t bytes = regions.expert_bytes_up; - const size_t file_off = regions.up_exps.offset + (size_t)expert_id * bytes; - if (file_off + bytes > mmap_size) { - if (err) *err = "up expert out of file bounds"; - return false; + runtime_->device_index.erase(cached); + } + + ++runtime_->device_cache_misses; + int victim = -1; + for (size_t i = 0; i < runtime_->device_slots.size(); ++i) { + const Runtime::DeviceSlot & slot = runtime_->device_slots[i]; + if (!slot.valid && !slot.pending && slot.compute_users == 0) { + victim = (int) i; + break; + } + } + if (victim < 0) { + uint64_t best_score = std::numeric_limits::max(); + for (size_t i = 0; i < runtime_->device_slots.size(); ++i) { + const Runtime::DeviceSlot & slot = runtime_->device_slots[i]; + if (!slot.valid || slot.pending || slot.compute_users != 0) continue; + const uint64_t age = runtime_->device_clock >= slot.last_touch + ? runtime_->device_clock - slot.last_touch : 0; + const uint64_t recency = age < 65535 ? 65535 - age : 0; + const uint64_t score = (slot.frequency << 16) | recency; + if (score < best_score) { + best_score = score; + victim = (int) i; } - std::memcpy(static_cast(pinned_buf_) + staging_offset, - file_base + file_off, bytes); - last_up_bytes_ = bytes; - staging_offset += bytes; } } + if (victim < 0) { + if (err) *err = "all SSD GPU expert-cache slots are busy"; + return false; + } + + const bool evicting = runtime_->device_slots[(size_t) victim].valid; + if (!stage_expert_async(layer, expert_id, victim, err)) return false; + Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) victim]; + slot.valid = true; + slot.cache_managed = true; + slot.key = {(int32_t) layer, (int32_t) expert_id}; + slot.frequency = 1; + slot.last_touch = ++runtime_->device_clock; + runtime_->device_index[key] = victim; + if (evicting) ++runtime_->device_cache_evictions; + *device_slot = victim; + return true; +} - // down - { - const size_t bytes = regions.expert_bytes_down; - const size_t file_off = regions.down_exps.offset + (size_t)expert_id * bytes; - if (file_off + bytes > mmap_size) { - if (err) *err = "down expert out of file bounds"; +bool MoeHybridStreamEngine::activate_device_slot(int device_slot, + std::string * err) { + if (!is_ready() || device_slot < 0 || + device_slot >= (int) runtime_->device_slots.size()) { + if (err) *err = "SSD device slot is out of range"; + return false; + } + ScopedGpuDevice device_scope(runtime_->device); + if (!device_scope.ready()) { + if (err) *err = "failed to select SSD stream GPU"; + return false; + } + Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) device_slot]; + if (slot.pending) { + const cudaError_t gpu_err = cudaEventSynchronize(slot.ready); + if (gpu_err != cudaSuccess) { + if (err) *err = std::string("expert H2D synchronization failed: ") + + cudaGetErrorString(gpu_err); return false; } - std::memcpy(static_cast(pinned_buf_) + staging_offset, - file_base + file_off, bytes); - last_down_bytes_ = bytes; - staging_offset += bytes; + slot.pending = false; + slot.host_lease.reset(); } - - if (staging_offset > scratch_size_) { - if (err) *err = "expert exceeds scratch buffer size"; + if (slot.layout.span_count < 2) { + if (err) *err = "SSD device slot has no complete expert"; return false; } + if (slot.cache_managed) { + if (slot.compute_users != 0) { + if (err) *err = "cached expert slot is already executing"; + return false; + } + ++slot.compute_users; + ++slot.frequency; + slot.last_touch = ++runtime_->device_clock; + } + runtime_->active_slot = device_slot; + return true; +} - // DMA pinned → GPU scratch (synchronous for now; async pipeline in eval function) - cudaError_t cuda_err = cudaMemcpy(gpu_scratch_, pinned_buf_, staging_offset, - cudaMemcpyHostToDevice); - if (cuda_err != cudaSuccess) { - if (err) *err = std::string("cudaMemcpy H2D failed: ") + cudaGetErrorString(cuda_err); - return false; +void MoeHybridStreamEngine::release_device_slot(int device_slot) { + if (!runtime_ || device_slot < 0 || + device_slot >= (int) runtime_->device_slots.size()) { + return; } + Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) device_slot]; + if (slot.cache_managed && slot.compute_users > 0) --slot.compute_users; + if (runtime_->active_slot == device_slot) runtime_->active_slot = -1; +} - // Set pointers into scratch - auto * scratch_bytes = static_cast(gpu_scratch_); - size_t off = 0; - if (regions.fused_gate_up) { - scratch_gate_ = scratch_bytes + off; - off += last_gate_bytes_; - scratch_up_ = nullptr; - } else { - scratch_gate_ = scratch_bytes + off; - off += last_gate_bytes_; - scratch_up_ = scratch_bytes + off; - off += last_up_bytes_; +int MoeHybridStreamEngine::device_slot_count() const { + return runtime_ ? (int) runtime_->device_slots.size() : 0; +} + +size_t MoeHybridStreamEngine::device_cache_bytes() const { + return runtime_ ? runtime_->device_pool_bytes : 0; +} + +ggml_backend_t MoeHybridStreamEngine::compute_backend() const { + return runtime_ ? runtime_->backend : nullptr; +} + +bool MoeHybridStreamEngine::stream_expert_sync(int layer, int expert_id, + std::string * err) { + if (!stage_expert_async(layer, expert_id, 0, err)) return false; + return activate_device_slot(0, err); +} + +bool MoeHybridStreamEngine::stream_expert_sync( + const void * mmap_data, size_t mmap_size, + const LayerExpertRegions & regions, int expert_id, + ggml_backend_t gpu_backend, std::string * err) { + (void) gpu_backend; + if (!is_ready()) { + if (err) *err = "stream engine is not initialized"; + return false; } - scratch_down_ = scratch_bytes + off; + if (!is_bound()) { + std::vector one_layer{regions}; + if (!runtime_->io->bind_source({mmap_data, mmap_size, -1}, one_layer, err)) return false; + } + return stream_expert_sync(0, expert_id, err); +} - return true; +const void * MoeHybridStreamEngine::scratch_gate_data() const { + if (!runtime_ || runtime_->active_slot < 0) return nullptr; + const Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) runtime_->active_slot]; + return static_cast(slot.data) + slot.layout.spans[0].device_offset; +} + +const void * MoeHybridStreamEngine::scratch_up_data() const { + if (!runtime_ || runtime_->active_slot < 0) return nullptr; + const Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) runtime_->active_slot]; + if (slot.layout.fused_gate_up) return nullptr; + return static_cast(slot.data) + slot.layout.spans[1].device_offset; +} + +const void * MoeHybridStreamEngine::scratch_down_data() const { + if (!runtime_ || runtime_->active_slot < 0) return nullptr; + const Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) runtime_->active_slot]; + const int index = slot.layout.fused_gate_up ? 1 : 2; + return static_cast(slot.data) + slot.layout.spans[index].device_offset; +} + +size_t MoeHybridStreamEngine::scratch_gate_bytes() const { + if (!runtime_ || runtime_->active_slot < 0) return 0; + return runtime_->device_slots[(size_t) runtime_->active_slot].layout.spans[0].bytes; +} + +size_t MoeHybridStreamEngine::scratch_up_bytes() const { + if (!runtime_ || runtime_->active_slot < 0) return 0; + const MoeExpertIoLayout & layout = + runtime_->device_slots[(size_t) runtime_->active_slot].layout; + return layout.fused_gate_up ? 0 : layout.spans[1].bytes; +} + +size_t MoeHybridStreamEngine::scratch_down_bytes() const { + if (!runtime_ || runtime_->active_slot < 0) return 0; + const MoeExpertIoLayout & layout = + runtime_->device_slots[(size_t) runtime_->active_slot].layout; + return layout.spans[layout.fused_gate_up ? 1 : 2].bytes; +} + +size_t MoeHybridStreamEngine::pinned_bytes() const { + return runtime_ && runtime_->io ? runtime_->io->total_host_bytes() : 0; +} + +size_t MoeHybridStreamEngine::scratch_bytes() const { + return runtime_ ? runtime_->device_pool_bytes : 0; +} + +const char * MoeHybridStreamEngine::io_backend_name() const { + return runtime_ && runtime_->io ? runtime_->io->effective_backend_name() : "uninitialized"; } -// ── Streaming prefill evaluation ──────────────────────────────────────────── +MoeNvmeStats MoeHybridStreamEngine::io_stats() const { + return runtime_ && runtime_->io ? runtime_->io->stats() : MoeNvmeStats{}; +} bool eval_moe_cold_experts_streaming( MoeHybridStreamEngine & engine, @@ -280,155 +772,199 @@ bool eval_moe_cold_experts_streaming( const float * selected_weights, int n_tokens, std::vector & out, - std::string * err) { + std::string * err, + int layer) { + + // The streamed tier can intentionally target a different owner from the + // caller (Strix for Lucebox, while hot/dense work stays on the R9700). + // Always build and launch the expert graph on the device that owns the + // stream slots. + if (engine.compute_backend()) gpu_backend = engine.compute_backend(); - const int n_embd = cfg.n_embd; + const int n_embd = cfg.expert_embd(); const int n_ff_exp = cfg.n_ff_exp; const int n_used = cfg.n_expert_used; const int total_slots = n_used * n_tokens; - out.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); + if (cfg.gated_activation == MoeGatedActivation::Situ && + (cfg.situ_beta <= 0.0f || cfg.situ_linear_beta <= 0.0f)) { + if (err) *err = "SiTU activation scales must be positive"; + return false; + } + + out.assign((size_t) n_embd * (size_t) n_tokens, 0.0f); if (!engine.is_ready()) { - if (err) *err = "stream engine not ready"; + if (err) *err = "stream engine is not ready"; return false; } - if (!mmap_data || mmap_size == 0) { - if (err) *err = "mmap not available"; + if (!engine.is_bound() && (!mmap_data || mmap_size == 0)) { + if (err) *err = "mmap is not available"; return false; } - // Identify unique cold experts needed across all tokens. - std::vector cold_needed((size_t)cfg.n_expert, false); + std::vector cold_needed((size_t) cfg.n_expert, false); for (int i = 0; i < total_slots; ++i) { const int32_t gid = selected_ids[i]; if (gid < 0 || gid >= cfg.n_expert) continue; - if (storage.hot_local_by_global[(size_t)gid] < 0) { - cold_needed[(size_t)gid] = true; - } + if (selected_weights[i] == 0.0f) continue; + if (storage.hot_local_by_global[(size_t) gid] < 0) cold_needed[(size_t) gid] = true; } std::vector unique_cold; - for (int e = 0; e < cfg.n_expert; ++e) { - if (cold_needed[(size_t)e]) unique_cold.push_back((int32_t)e); + for (int expert = 0; expert < cfg.n_expert; ++expert) { + if (cold_needed[(size_t) expert]) unique_cold.push_back((int32_t) expert); } - if (unique_cold.empty()) return true; - // Prefetch all cold experts via madvise - engine.prefetch_cold_experts(mmap_data, mmap_size, regions, unique_cold.data(), (int)unique_cold.size()); + const bool cache_pipeline = engine.is_bound(); - // For each unique cold expert: stream to GPU, compute ALL tokens that selected it - // in a single batched matmul graph. - for (int32_t cold_eid : unique_cold) { - // Stream expert weights to GPU scratch - if (!engine.stream_expert_sync(mmap_data, mmap_size, regions, cold_eid, gpu_backend, err)) { - return false; - } + // Admit every actual route before compute. io_uring sees the whole batch; + // the thread fallback obtains enough outstanding reads to saturate NVMe. + if (cache_pipeline) { + engine.request_experts(layer, unique_cold.data(), (int) unique_cold.size(), + MoeNvmePriority::Demand); + } + + int staged_device_slot = 0; + if (cache_pipeline) { + if (!engine.stage_expert_cached_async( + layer, unique_cold[0], &staged_device_slot, err)) return false; + } else { + if (!engine.stream_expert_sync(mmap_data, mmap_size, regions, + unique_cold[0], gpu_backend, err)) return false; + } + + for (size_t cold_index = 0; cold_index < unique_cold.size(); ++cold_index) { + const int32_t cold_eid = unique_cold[cold_index]; + const int current_device_slot = staged_device_slot; + if (cache_pipeline && + !engine.activate_device_slot(current_device_slot, err)) return false; + auto release_current = [&]() { + if (cache_pipeline) engine.release_device_slot(current_device_slot); + }; - // Gather all tokens that selected this expert - struct TokenHit { int ti; float weight; }; + struct TokenHit { int token; float weight; }; std::vector hits; - hits.reserve((size_t)n_tokens); - for (int ti = 0; ti < n_tokens; ++ti) { + hits.reserve((size_t) n_tokens); + for (int token = 0; token < n_tokens; ++token) { for (int k = 0; k < n_used; ++k) { - const int slot = ti * n_used + k; + const int slot = token * n_used + k; if (selected_ids[slot] != cold_eid) continue; - const float w = selected_weights[slot]; - if (w != 0.0f) hits.push_back({ti, w}); - break; // each expert selected at most once per token + if (selected_weights[slot] != 0.0f) { + hits.push_back({token, selected_weights[slot]}); + } + break; } } - if (hits.empty()) continue; - - const int batch = (int)hits.size(); + if (hits.empty()) { + release_current(); + continue; + } - // Build batched input: [n_embd, batch] - std::vector batch_input((size_t)n_embd * (size_t)batch); + const int batch = (int) hits.size(); + std::vector batch_input((size_t) n_embd * (size_t) batch); for (int i = 0; i < batch; ++i) { - const float * src = cur_host + (size_t)hits[(size_t)i].ti * (size_t)n_embd; - std::memcpy(batch_input.data() + (size_t)i * (size_t)n_embd, src, sizeof(float) * (size_t)n_embd); + const float * src = cur_host + (size_t) hits[(size_t) i].token * (size_t) n_embd; + std::memcpy(batch_input.data() + (size_t) i * (size_t) n_embd, + src, sizeof(float) * (size_t) n_embd); } - // Build single ggml graph for this expert with all tokens batched ggml_init_params ip{}; ip.mem_size = 32 * 1024 * 1024; ip.mem_buffer = nullptr; ip.no_alloc = true; ggml_context * ctx = ggml_init(ip); if (!ctx) { - if (err) *err = "ggml_init failed in streaming eval"; + if (err) *err = "ggml_init failed in SSD streaming eval"; + release_current(); return false; } ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, batch); ggml_set_input(inp); - - // Weight tensors pointing into GPU scratch (same for all tokens in batch) ggml_tensor * gate_t = nullptr; ggml_tensor * up_t = nullptr; ggml_tensor * down_t = nullptr; ggml_tensor * gate_up_t = nullptr; - if (regions.fused_gate_up) { - gate_up_t = ggml_new_tensor_2d(ctx, desc.ffn_gate_up_exps->type, n_embd, 2 * n_ff_exp); + gate_up_t = ggml_new_tensor_2d(ctx, desc.ffn_gate_up_exps->type, + n_embd, 2 * n_ff_exp); + down_t = ggml_new_tensor_2d(ctx, desc.ffn_down_exps->type, + n_ff_exp, n_embd); ggml_set_input(gate_up_t); - down_t = ggml_new_tensor_2d(ctx, desc.ffn_down_exps->type, n_ff_exp, n_embd); ggml_set_input(down_t); } else { - gate_t = ggml_new_tensor_2d(ctx, desc.ffn_gate_exps->type, n_embd, n_ff_exp); + gate_t = ggml_new_tensor_2d(ctx, desc.ffn_gate_exps->type, + n_embd, n_ff_exp); + up_t = ggml_new_tensor_2d(ctx, desc.ffn_up_exps->type, + n_embd, n_ff_exp); + down_t = ggml_new_tensor_2d(ctx, desc.ffn_down_exps->type, + n_ff_exp, n_embd); ggml_set_input(gate_t); - up_t = ggml_new_tensor_2d(ctx, desc.ffn_up_exps->type, n_embd, n_ff_exp); ggml_set_input(up_t); - down_t = ggml_new_tensor_2d(ctx, desc.ffn_down_exps->type, n_ff_exp, n_embd); ggml_set_input(down_t); } - // FFN graph: out = down(silu(gate(x)) * up(x)) — batched over all tokens - ggml_tensor * gu = nullptr; + auto apply_gated_activation = [&](ggml_tensor * gate, + ggml_tensor * up) -> ggml_tensor * { + if (cfg.gated_activation == MoeGatedActivation::Situ) { + ggml_tensor * nonlinear = ggml_scale(ctx, gate, 1.0f / cfg.situ_beta); + nonlinear = ggml_tanh(ctx, nonlinear); + nonlinear = ggml_scale(ctx, nonlinear, cfg.situ_beta); + nonlinear = ggml_mul(ctx, nonlinear, ggml_sigmoid(ctx, gate)); + ggml_tensor * linear = ggml_scale( + ctx, up, 1.0f / cfg.situ_linear_beta); + linear = ggml_tanh(ctx, linear); + linear = ggml_scale(ctx, linear, cfg.situ_linear_beta); + return ggml_mul(ctx, nonlinear, linear); + } + if (cfg.swiglu_clamp > 0.0f) { + return ggml_swiglu_ds4_split(ctx, gate, up, cfg.swiglu_clamp); + } + return ggml_swiglu_split(ctx, gate, up); + }; + + ggml_tensor * gated = nullptr; if (gate_up_t) { - ggml_tensor * gate_up_out = ggml_mul_mat(ctx, gate_up_t, inp); // [2*n_ff, batch] - if (desc.ffn_gate_up_exps_s != 1.0f) + ggml_tensor * gate_up_out = ggml_mul_mat(ctx, gate_up_t, inp); + if (desc.ffn_gate_up_exps_s != 1.0f) { gate_up_out = ggml_scale(ctx, gate_up_out, desc.ffn_gate_up_exps_s); - ggml_tensor * g_part = ggml_view_2d(ctx, gate_up_out, n_ff_exp, batch, - gate_up_out->nb[1], 0); - ggml_tensor * u_part = ggml_view_2d(ctx, gate_up_out, n_ff_exp, batch, - gate_up_out->nb[1], - (size_t)n_ff_exp * sizeof(float)); - g_part = ggml_cont(ctx, g_part); - u_part = ggml_cont(ctx, u_part); - gu = ggml_swiglu_split(ctx, g_part, u_part); + } + ggml_tensor * gate_part = ggml_view_2d( + ctx, gate_up_out, n_ff_exp, batch, gate_up_out->nb[1], 0); + ggml_tensor * up_part = ggml_view_2d( + ctx, gate_up_out, n_ff_exp, batch, gate_up_out->nb[1], + (size_t) n_ff_exp * sizeof(float)); + gate_part = ggml_cont(ctx, gate_part); + up_part = ggml_cont(ctx, up_part); + gated = apply_gated_activation(gate_part, up_part); } else { - ggml_tensor * g = ggml_mul_mat(ctx, gate_t, inp); // [n_ff, batch] - if (desc.ffn_gate_exps_s != 1.0f) - g = ggml_scale(ctx, g, desc.ffn_gate_exps_s); - ggml_tensor * u = ggml_mul_mat(ctx, up_t, inp); // [n_ff, batch] - if (desc.ffn_up_exps_s != 1.0f) - u = ggml_scale(ctx, u, desc.ffn_up_exps_s); - gu = ggml_swiglu_split(ctx, g, u); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_t, inp); + if (desc.ffn_gate_exps_s != 1.0f) gate = ggml_scale(ctx, gate, desc.ffn_gate_exps_s); + ggml_tensor * up = ggml_mul_mat(ctx, up_t, inp); + if (desc.ffn_up_exps_s != 1.0f) up = ggml_scale(ctx, up, desc.ffn_up_exps_s); + gated = apply_gated_activation(gate, up); } - - ggml_tensor * expert_out = ggml_mul_mat(ctx, down_t, gu); // [n_embd, batch] - if (desc.ffn_down_exps_s != 1.0f) + ggml_tensor * expert_out = ggml_mul_mat(ctx, down_t, gated); + if (desc.ffn_down_exps_s != 1.0f) { expert_out = ggml_scale(ctx, expert_out, desc.ffn_down_exps_s); + } - ggml_cgraph * gf = ggml_new_graph_custom(ctx, 512, false); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 512, false); ggml_set_output(expert_out); - ggml_build_forward_expand(gf, expert_out); - - ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(gpu_backend)); - if (!ggml_gallocr_alloc_graph(alloc, gf)) { - if (err) *err = "streaming eval gallocr failed"; - ggml_gallocr_free(alloc); + ggml_build_forward_expand(graph, expert_out); + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(gpu_backend)); + if (!alloc || !ggml_gallocr_alloc_graph(alloc, graph)) { + if (err) *err = "SSD streaming eval graph allocation failed"; + if (alloc) ggml_gallocr_free(alloc); ggml_free(ctx); + release_current(); return false; } - // Upload batched input (host → GPU via ggml_backend_tensor_set) - ggml_backend_tensor_set(inp, batch_input.data(), 0, sizeof(float) * (size_t)n_embd * (size_t)batch); - - // Point weight tensors directly at GPU scratch (device-to-device, no copy needed). - // The gallocr allocated these tensors on the same GPU, but we override their data - // pointers to point at our pre-loaded scratch buffer. + ggml_backend_tensor_set(inp, batch_input.data(), 0, + sizeof(float) * (size_t) n_embd * (size_t) batch); if (gate_up_t) { gate_up_t->data = const_cast(engine.scratch_gate_data()); down_t->data = const_cast(engine.scratch_down_data()); @@ -438,33 +974,51 @@ bool eval_moe_cold_experts_streaming( down_t->data = const_cast(engine.scratch_down_data()); } - auto st = ggml_backend_graph_compute(gpu_backend, gf); - if (st != GGML_STATUS_SUCCESS) { - if (err) *err = "streaming eval compute failed"; + const ggml_status status = ggml_backend_graph_compute_async(gpu_backend, graph); + if (status != GGML_STATUS_SUCCESS) { + if (err) *err = "SSD streaming expert compute launch failed"; ggml_gallocr_free(alloc); ggml_free(ctx); + release_current(); return false; } - // Read batched result [n_embd, batch] and scatter-accumulate with weights - std::vector batch_result((size_t)n_embd * (size_t)batch); - ggml_backend_tensor_get(expert_out, batch_result.data(), 0, sizeof(float) * (size_t)n_embd * (size_t)batch); - - for (int i = 0; i < batch; ++i) { - const float w = hits[(size_t)i].weight; - const int ti = hits[(size_t)i].ti; - const float * res = batch_result.data() + (size_t)i * (size_t)n_embd; - float * out_tok = out.data() + (size_t)ti * (size_t)n_embd; - for (int j = 0; j < n_embd; ++j) { - out_tok[j] += w * res[(size_t)j]; + // Compute N is now running. Wait for the already-issued disk read of + // N+1 and enqueue its H2D into a different device slot. + if (cold_index + 1 < unique_cold.size() && cache_pipeline) { + int next_slot = -1; + if (!engine.stage_expert_cached_async( + layer, unique_cold[cold_index + 1], &next_slot, err)) { + ggml_backend_synchronize(gpu_backend); + release_current(); + ggml_gallocr_free(alloc); + ggml_free(ctx); + return false; } + staged_device_slot = next_slot; } + ggml_backend_synchronize(gpu_backend); + std::vector batch_result((size_t) n_embd * (size_t) batch); + ggml_backend_tensor_get(expert_out, batch_result.data(), 0, + sizeof(float) * (size_t) n_embd * (size_t) batch); + for (int i = 0; i < batch; ++i) { + const float weight = hits[(size_t) i].weight; + float * dst = out.data() + (size_t) hits[(size_t) i].token * (size_t) n_embd; + const float * src = batch_result.data() + (size_t) i * (size_t) n_embd; + for (int j = 0; j < n_embd; ++j) dst[j] += weight * src[(size_t) j]; + } ggml_gallocr_free(alloc); ggml_free(ctx); - } + release_current(); + if (cold_index + 1 < unique_cold.size() && !cache_pipeline) { + if (!engine.stream_expert_sync(mmap_data, mmap_size, regions, + unique_cold[cold_index + 1], + gpu_backend, err)) return false; + } + } return true; } -} // namespace dflash::common +} // namespace dflash::common diff --git a/server/src/common/moe_hybrid_stream.h b/server/src/common/moe_hybrid_stream.h index d5e929ae9..0d1c1f73c 100644 --- a/server/src/common/moe_hybrid_stream.h +++ b/server/src/common/moe_hybrid_stream.h @@ -1,34 +1,42 @@ -// MoE hybrid prefill streaming engine — DMA pipeline for cold expert offload. +// Heterogeneous MoE SSD execution tier. // -// Streams cold expert weight slices from mmap (page cache) through a pinned -// host staging buffer to GPU scratch memory, pipelined with GPU compute on -// previously-transferred experts. Used during prefill when T >= threshold. +// Exact routed experts move through a bounded NVMe -> page-locked host -> GPU +// pipeline. The model storage format remains separate: this runtime only +// consumes model-neutral LayerExpertRegions produced by a loader. #pragma once #include "moe_hybrid_types.h" #include "moe_hybrid_storage.h" +#include "moe_nvme_scheduler.h" #include "ggml.h" #include "ggml-backend.h" #include #include +#include #include #include namespace dflash::common { -// Configuration for the stream engine. struct MoeStreamConfig { - int prefill_threshold = 8; // min n_tokens to activate streaming - int prefetch_layers = 2; // how many layers ahead to madvise + int prefill_threshold = 8; + int prefetch_layers = 2; + int device_slots = 2; // double buffering is the minimum useful pipeline + // Optional adaptive GPU expert-cache budget. Zero keeps only the pipeline + // slots. The hardware planner can safely assign otherwise-unused Strix + // memory here while retaining its KV/graph reserve. + size_t device_cache_bytes = 0; + MoeNvmeConfig nvme{}; + + static MoeStreamConfig from_env(); }; -// Streaming engine: manages pinned staging buffer, GPU scratch, and DMA pipeline. class MoeHybridStreamEngine { public: - MoeHybridStreamEngine() = default; + MoeHybridStreamEngine(); ~MoeHybridStreamEngine(); MoeHybridStreamEngine(const MoeHybridStreamEngine &) = delete; @@ -36,63 +44,85 @@ class MoeHybridStreamEngine { MoeHybridStreamEngine(MoeHybridStreamEngine &&) noexcept; MoeHybridStreamEngine & operator=(MoeHybridStreamEngine &&) noexcept; - // Initialize the engine with a maximum expert size (bytes for one expert's - // gate+up+down tensors). Allocates pinned host buffer and GPU scratch. - bool init(ggml_backend_t gpu_backend, size_t max_expert_bytes, std::string * err = nullptr); + // Compatibility initialization for synthetic callers. Production callers + // should use the storage overload so actual file reads (and io_uring) are + // available instead of relying only on mmap page faults. + bool init(ggml_backend_t gpu_backend, size_t max_expert_bytes, + std::string * err = nullptr); + bool init(ggml_backend_t gpu_backend, size_t max_expert_bytes, + const MoeHybridStorage & storage, + std::string * err = nullptr); + bool init(ggml_backend_t gpu_backend, size_t max_expert_bytes, + const MoeHybridStorage & storage, + const MoeStreamConfig & config, + std::string * err = nullptr); + + bool bind_storage(const MoeHybridStorage & storage, std::string * err = nullptr); + bool bind_sources(const std::vector & sources, + const std::vector & layer_regions, + std::string * err = nullptr); bool is_ready() const; + bool is_bound() const; void destroy(); - // Issue madvise(WILLNEED) for the specified cold experts in the given layer. - // Call this as early as possible (e.g. at start of layer or N layers ahead). + // Queue exact experts without waiting. Demand requests always outrank + // speculative prefetches and can cancel queued speculation. + void request_experts(int layer, const int32_t * expert_ids, int count, + MoeNvmePriority priority = MoeNvmePriority::Prefetch); + + // Compatibility page-cache hint. New code should call request_experts(), + // which performs real asynchronous reads into the bounded host cache. void prefetch_cold_experts(const void * mmap_data, size_t mmap_size, const LayerExpertRegions & regions, const int32_t * cold_expert_ids, int n_cold); - // Stream a single cold expert from mmap to GPU scratch and return a ggml - // tensor view over the scratch memory for each weight matrix. - // This is a BLOCKING operation (synchronous DMA). For pipelined usage, - // use the async variants below. + // Queue one H2D transfer into a device slot, then activate it after its + // completion event. Different slots allow transfer N+1 to overlap compute N. + bool stage_expert_async(int layer, int expert_id, int device_slot, + std::string * err = nullptr); + + // Cache-aware form used by production inference. On a hit it returns the + // existing Strix slot without host or SSD traffic. On a miss it selects an + // unpinned LFRU victim and starts the same asynchronous upload pipeline. + bool stage_expert_cached_async(int layer, int expert_id, int * device_slot, + std::string * err = nullptr); + bool activate_device_slot(int device_slot, std::string * err = nullptr); + void release_device_slot(int device_slot); + int device_slot_count() const; + size_t device_cache_bytes() const; + ggml_backend_t compute_backend() const; + + bool stream_expert_sync(int layer, int expert_id, + std::string * err = nullptr); + + // Legacy form: lazily binds a single synthetic layer. bool stream_expert_sync(const void * mmap_data, size_t mmap_size, const LayerExpertRegions & regions, int expert_id, ggml_backend_t gpu_backend, std::string * err = nullptr); - // Get tensor pointers into the GPU scratch buffer after a successful stream. - // Valid until next stream call. Tensors are transient (not owned by any context). - const void * scratch_gate_data() const { return scratch_gate_; } - const void * scratch_up_data() const { return scratch_up_; } - const void * scratch_down_data() const { return scratch_down_; } - size_t scratch_gate_bytes() const { return last_gate_bytes_; } - size_t scratch_up_bytes() const { return last_up_bytes_; } - size_t scratch_down_bytes() const { return last_down_bytes_; } + const void * scratch_gate_data() const; + const void * scratch_up_data() const; + const void * scratch_down_data() const; + size_t scratch_gate_bytes() const; + size_t scratch_up_bytes() const; + size_t scratch_down_bytes() const; - // Total pinned buffer size. - size_t pinned_bytes() const { return pinned_size_; } - // Total GPU scratch size. - size_t scratch_bytes() const { return scratch_size_; } + size_t pinned_bytes() const; + size_t scratch_bytes() const; + const char * io_backend_name() const; + MoeNvmeStats io_stats() const; private: - void * pinned_buf_ = nullptr; // cudaMallocHost'd staging buffer - size_t pinned_size_ = 0; - - void * gpu_scratch_ = nullptr; // GPU device memory for one expert - size_t scratch_size_ = 0; - ggml_backend_t backend_ = nullptr; - - // Offsets into scratch for last-streamed expert - void * scratch_gate_ = nullptr; - void * scratch_up_ = nullptr; - void * scratch_down_ = nullptr; - size_t last_gate_bytes_ = 0; - size_t last_up_bytes_ = 0; - size_t last_down_bytes_ = 0; + struct Runtime; + std::unique_ptr runtime_; }; -// Evaluate cold experts by streaming from mmap to GPU, pipelined. -// Hot experts are already computed (result in hot_partial). -// Returns combined cold expert contribution in out (sized n_embd * n_tokens). +// Evaluate the cold contribution for one layer. All routed SSD requests are +// admitted before compute starts, then double-buffered H2D runs concurrently +// with the preceding expert graph. bool eval_moe_cold_experts_streaming( MoeHybridStreamEngine & engine, ggml_backend_t gpu_backend, @@ -107,6 +137,7 @@ bool eval_moe_cold_experts_streaming( const float * selected_weights, int n_tokens, std::vector & out, - std::string * err = nullptr); + std::string * err = nullptr, + int layer = 0); -} // namespace dflash::common +} // namespace dflash::common diff --git a/server/src/common/moe_hybrid_types.h b/server/src/common/moe_hybrid_types.h index bdb22e47e..fbebca6b4 100644 --- a/server/src/common/moe_hybrid_types.h +++ b/server/src/common/moe_hybrid_types.h @@ -23,10 +23,18 @@ enum class MoeHybridColdBackend { Gpu, }; +enum class MoeGatedActivation { + SwiGlu, + Situ, +}; + // ─── MoE architecture config (model-agnostic) ────────────────────────── struct MoeHybridConfig { int n_embd = 0; // hidden dimension + // Some MoEs project hidden states into a smaller routed-expert latent + // space (Kimi K3: 7168 -> 3584). Zero means use n_embd. + int n_expert_embd = 0; int n_expert = 0; // total experts per layer int n_expert_used = 0; // top-k selected per token int n_ff_exp = 0; // routed expert intermediate dimension @@ -34,6 +42,9 @@ struct MoeHybridConfig { int n_layer = 0; // number of MoE layers int first_moe_layer = 0; // index of first MoE layer (e.g., 0 for qwen35moe, 1 for laguna) float swiglu_clamp = 0.0f; // 0 = regular SwiGLU; >0 clamps gate upper/up symmetric (DS4) + MoeGatedActivation gated_activation = MoeGatedActivation::SwiGlu; + float situ_beta = 4.0f; + float situ_linear_beta = 25.0f; MoeHybridColdBackend cold_expert_backend = MoeHybridColdBackend::Cpu; bool materialize_hot_experts = true; bool materialize_cold_experts = true; @@ -43,6 +54,8 @@ struct MoeHybridConfig { // On sm_75 (Turing) and gfx1151, the kernel has illegal memory accesses // with reduced stacks, requiring the <=4-token sub-batch workaround. bool mmq_safe_full_batch = false; + + int expert_embd() const { return n_expert_embd > 0 ? n_expert_embd : n_embd; } }; // ─── Per-layer expert tensor descriptor ───────────────────────────────── diff --git a/server/src/common/moe_nvme_scheduler.cpp b/server/src/common/moe_nvme_scheduler.cpp new file mode 100644 index 000000000..bfc481b9a --- /dev/null +++ b/server/src/common/moe_nvme_scheduler.cpp @@ -0,0 +1,1636 @@ +// Model-neutral asynchronous SSD scheduler for routed MoE weights. + +#include "moe_nvme_scheduler.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#else +#include +#include +#include +#include +#include +#endif + +#if defined(__linux__) +#include +#include +#include +#endif + +namespace dflash::common { +namespace { + +using Clock = std::chrono::steady_clock; + +uint64_t elapsed_ns(Clock::time_point begin, Clock::time_point end) { + return (uint64_t) std::chrono::duration_cast(end - begin).count(); +} + +bool checked_add(size_t a, size_t b, size_t & out) { + if (a > std::numeric_limits::max() - b) return false; + out = a + b; + return true; +} + +bool checked_mul(size_t a, size_t b, size_t & out) { + if (a != 0 && b > std::numeric_limits::max() / a) return false; + out = a * b; + return true; +} + +bool range_in_bounds(size_t offset, size_t length, size_t total) { + return offset <= total && length <= total - offset; +} + +bool is_power_of_two(size_t value) { + return value != 0 && (value & (value - 1)) == 0; +} + +bool align_up_checked(size_t value, size_t alignment, size_t & out) { + if (!is_power_of_two(alignment)) return false; + const size_t mask = alignment - 1; + if (value > std::numeric_limits::max() - mask) return false; + out = (value + mask) & ~mask; + return true; +} + +size_t align_down(size_t value, size_t alignment) { + return value & ~(alignment - 1); +} + +int parse_bounded_int(const char * name, int current, int lo, int hi) { + const char * value = std::getenv(name); + if (!value || !value[0]) return current; + char * end = nullptr; + errno = 0; + const long parsed = std::strtol(value, &end, 10); + if (errno != 0 || end == value || *end != '\0' || parsed < lo || parsed > hi) { + std::fprintf(stderr, "[moe-nvme] ignoring invalid %s=%s (want %d..%d)\n", + name, value, lo, hi); + return current; + } + return (int) parsed; +} + +std::string lowercase(const char * value) { + std::string out = value ? value : ""; + for (char & ch : out) { + if (ch >= 'A' && ch <= 'Z') ch = (char) (ch - 'A' + 'a'); + } + return out; +} + +struct KeyHash { + size_t operator()(const MoeExpertKey & key) const noexcept { + const uint64_t a = (uint32_t) key.layer; + const uint64_t b = (uint32_t) key.expert; + uint64_t x = (a << 32) | b; + x ^= x >> 30; + x *= UINT64_C(0xbf58476d1ce4e5b9); + x ^= x >> 27; + x *= UINT64_C(0x94d049bb133111eb); + x ^= x >> 31; + return (size_t) x; + } +}; + +uint64_t physical_memory_bytes() { +#if defined(_WIN32) + return 0; +#else + const long pages = ::sysconf(_SC_PHYS_PAGES); + const long page_size = ::sysconf(_SC_PAGESIZE); + if (pages <= 0 || page_size <= 0) return 0; + const uint64_t p = (uint64_t) pages; + const uint64_t s = (uint64_t) page_size; + if (p > std::numeric_limits::max() / s) return 0; + return p * s; +#endif +} + +#if !defined(_WIN32) +bool pread_full(int fd, uint8_t * dst, size_t bytes, size_t offset, std::string & err) { + size_t done = 0; + while (done < bytes) { + const size_t remaining = bytes - done; + const size_t chunk = std::min(remaining, (size_t) std::numeric_limits::max()); + const ssize_t got = ::pread(fd, dst + done, chunk, (off_t) (offset + done)); + if (got < 0) { + if (errno == EINTR) continue; + err = std::string("pread failed: ") + std::strerror(errno); + return false; + } + if (got == 0) { + err = "short read at end of model file"; + return false; + } + done += (size_t) got; + } + return true; +} +#endif + +#if defined(__linux__) + +// Small dependency-free io_uring wrapper. The ABI is Linux UAPI, so the +// inference binary does not need liburing. The implementation follows the +// kernel io_uring interface and uses conservative flags for old enterprise +// kernels. Files and page-locked slots are registered when the kernel accepts +// them, reducing per-read pinning and fd-table overhead. +class RawIoUring { +public: + ~RawIoUring() { close(); } + + bool open(unsigned entries, const std::vector & active_fds, + const std::vector & buffers, size_t buffer_bytes, + std::string & err) { + close(); + std::memset(¶ms_, 0, sizeof(params_)); + fd_ = (int) ::syscall(SYS_io_uring_setup, entries, ¶ms_); + if (fd_ < 0) { + err = std::string("io_uring_setup failed: ") + std::strerror(errno); + return false; + } + + sq_ring_bytes_ = params_.sq_off.array + params_.sq_entries * sizeof(unsigned); + cq_ring_bytes_ = params_.cq_off.cqes + params_.cq_entries * sizeof(io_uring_cqe); + if (params_.features & IORING_FEAT_SINGLE_MMAP) { + const size_t both = std::max(sq_ring_bytes_, cq_ring_bytes_); + sq_ring_ = ::mmap(nullptr, both, PROT_READ | PROT_WRITE, + MAP_SHARED | MAP_POPULATE, fd_, IORING_OFF_SQ_RING); + if (sq_ring_ == MAP_FAILED) { + sq_ring_ = nullptr; + err = std::string("io_uring SQ/CQ mmap failed: ") + std::strerror(errno); + close(); + return false; + } + cq_ring_ = sq_ring_; + sq_ring_bytes_ = both; + cq_ring_bytes_ = both; + single_mmap_ = true; + } else { + sq_ring_ = ::mmap(nullptr, sq_ring_bytes_, PROT_READ | PROT_WRITE, + MAP_SHARED | MAP_POPULATE, fd_, IORING_OFF_SQ_RING); + cq_ring_ = ::mmap(nullptr, cq_ring_bytes_, PROT_READ | PROT_WRITE, + MAP_SHARED | MAP_POPULATE, fd_, IORING_OFF_CQ_RING); + if (sq_ring_ == MAP_FAILED || cq_ring_ == MAP_FAILED) { + if (sq_ring_ == MAP_FAILED) sq_ring_ = nullptr; + if (cq_ring_ == MAP_FAILED) cq_ring_ = nullptr; + err = std::string("io_uring ring mmap failed: ") + std::strerror(errno); + close(); + return false; + } + } + + sqes_bytes_ = params_.sq_entries * sizeof(io_uring_sqe); + sqes_ = static_cast(::mmap( + nullptr, sqes_bytes_, PROT_READ | PROT_WRITE, + MAP_SHARED | MAP_POPULATE, fd_, IORING_OFF_SQES)); + if (sqes_ == MAP_FAILED) { + sqes_ = nullptr; + err = std::string("io_uring SQE mmap failed: ") + std::strerror(errno); + close(); + return false; + } + + auto * sq = static_cast(sq_ring_); + auto * cq = static_cast(cq_ring_); + sq_head_ = reinterpret_cast(sq + params_.sq_off.head); + sq_tail_ = reinterpret_cast(sq + params_.sq_off.tail); + sq_mask_ = reinterpret_cast(sq + params_.sq_off.ring_mask); + sq_entries_ = reinterpret_cast(sq + params_.sq_off.ring_entries); + sq_array_ = reinterpret_cast(sq + params_.sq_off.array); + cq_head_ = reinterpret_cast(cq + params_.cq_off.head); + cq_tail_ = reinterpret_cast(cq + params_.cq_off.tail); + cq_mask_ = reinterpret_cast(cq + params_.cq_off.ring_mask); + cqes_ = reinterpret_cast(cq + params_.cq_off.cqes); + sqe_head_ = sqe_tail_ = 0; + + if (!active_fds.empty() && + ::syscall(SYS_io_uring_register, fd_, IORING_REGISTER_FILES, + active_fds.data(), (unsigned) active_fds.size()) == 0) { + fixed_file_ = true; + } + + iovecs_.resize(buffers.size()); + for (size_t i = 0; i < buffers.size(); ++i) { + iovecs_[i].iov_base = buffers[i]; + iovecs_[i].iov_len = buffer_bytes; + } + if (!iovecs_.empty() && + ::syscall(SYS_io_uring_register, fd_, IORING_REGISTER_BUFFERS, + iovecs_.data(), (unsigned) iovecs_.size()) == 0) { + fixed_buffers_ = true; + } + return true; + } + + void close() { + if (fd_ >= 0 && fixed_buffers_) { + (void) ::syscall(SYS_io_uring_register, fd_, IORING_UNREGISTER_BUFFERS, + nullptr, 0U); + } + if (fd_ >= 0 && fixed_file_) { + (void) ::syscall(SYS_io_uring_register, fd_, IORING_UNREGISTER_FILES, + nullptr, 0U); + } + fixed_buffers_ = false; + fixed_file_ = false; + iovecs_.clear(); + if (sqes_) ::munmap(sqes_, sqes_bytes_); + sqes_ = nullptr; + if (single_mmap_) { + if (sq_ring_) ::munmap(sq_ring_, sq_ring_bytes_); + } else { + if (sq_ring_) ::munmap(sq_ring_, sq_ring_bytes_); + if (cq_ring_) ::munmap(cq_ring_, cq_ring_bytes_); + } + sq_ring_ = cq_ring_ = nullptr; + if (fd_ >= 0) ::close(fd_); + fd_ = -1; + single_mmap_ = false; + sq_ring_bytes_ = cq_ring_bytes_ = sqes_bytes_ = 0; + sq_head_ = sq_tail_ = sq_mask_ = sq_entries_ = sq_array_ = nullptr; + cq_head_ = cq_tail_ = cq_mask_ = nullptr; + cqes_ = nullptr; + sqe_head_ = sqe_tail_ = 0; + } + + io_uring_sqe * get_sqe() { + const unsigned kernel_head = __atomic_load_n(sq_head_, __ATOMIC_ACQUIRE); + if (sqe_tail_ - kernel_head >= *sq_entries_) return nullptr; + io_uring_sqe * sqe = &sqes_[sqe_tail_ & *sq_mask_]; + std::memset(sqe, 0, sizeof(*sqe)); + ++sqe_tail_; + return sqe; + } + + void prepare_read(io_uring_sqe * sqe, + const std::vector & active_fds, + uint32_t source_index, int slot, + void * dst, uint32_t bytes, uint64_t offset, + uint64_t user_data, bool force_async) const { + sqe->opcode = fixed_buffers_ ? IORING_OP_READ_FIXED : IORING_OP_READ; + sqe->fd = fixed_file_ ? (int) source_index : active_fds[source_index]; + sqe->off = offset; + sqe->addr = (uint64_t) (uintptr_t) dst; + sqe->len = bytes; + sqe->user_data = user_data; + if (fixed_buffers_) sqe->buf_index = (uint16_t) slot; + if (fixed_file_) sqe->flags |= IOSQE_FIXED_FILE; + if (force_async) sqe->flags |= IOSQE_ASYNC; + } + + bool submit_all(std::string & err) { + unsigned kernel_tail = __atomic_load_n(sq_tail_, __ATOMIC_RELAXED); + const unsigned mask = *sq_mask_; + const unsigned count = sqe_tail_ - sqe_head_; + for (unsigned i = 0; i < count; ++i) { + sq_array_[kernel_tail & mask] = sqe_head_ & mask; + ++kernel_tail; + ++sqe_head_; + } + __atomic_store_n(sq_tail_, kernel_tail, __ATOMIC_RELEASE); + + while (__atomic_load_n(sq_head_, __ATOMIC_ACQUIRE) != kernel_tail) { + const unsigned pending = kernel_tail - __atomic_load_n(sq_head_, __ATOMIC_ACQUIRE); + const int rc = (int) ::syscall(SYS_io_uring_enter, fd_, pending, 0U, 0U, nullptr, 0U); + if (rc >= 0) continue; + if (errno == EINTR) continue; + err = std::string("io_uring_enter submit failed: ") + std::strerror(errno); + return false; + } + return true; + } + + bool wait_cqe(io_uring_cqe & out, std::string & err) { + for (;;) { + const unsigned head = __atomic_load_n(cq_head_, __ATOMIC_RELAXED); + const unsigned tail = __atomic_load_n(cq_tail_, __ATOMIC_ACQUIRE); + if (head != tail) { + out = cqes_[head & *cq_mask_]; + __atomic_store_n(cq_head_, head + 1, __ATOMIC_RELEASE); + return true; + } + const int rc = (int) ::syscall(SYS_io_uring_enter, fd_, 0U, 1U, + IORING_ENTER_GETEVENTS, nullptr, 0U); + if (rc >= 0) continue; + if (errno == EINTR) continue; + err = std::string("io_uring_enter wait failed: ") + std::strerror(errno); + return false; + } + } + +private: + int fd_ = -1; + io_uring_params params_{}; + void * sq_ring_ = nullptr; + void * cq_ring_ = nullptr; + io_uring_sqe * sqes_ = nullptr; + size_t sq_ring_bytes_ = 0; + size_t cq_ring_bytes_ = 0; + size_t sqes_bytes_ = 0; + bool single_mmap_ = false; + bool fixed_file_ = false; + bool fixed_buffers_ = false; + std::vector iovecs_; + + unsigned * sq_head_ = nullptr; + unsigned * sq_tail_ = nullptr; + unsigned * sq_mask_ = nullptr; + unsigned * sq_entries_ = nullptr; + unsigned * sq_array_ = nullptr; + unsigned * cq_head_ = nullptr; + unsigned * cq_tail_ = nullptr; + unsigned * cq_mask_ = nullptr; + io_uring_cqe * cqes_ = nullptr; + unsigned sqe_head_ = 0; + unsigned sqe_tail_ = 0; +}; + +#endif // __linux__ + +} // namespace + +MoeNvmeConfig MoeNvmeConfig::from_env() { + return from_env(MoeNvmeConfig{}); +} + +MoeNvmeConfig MoeNvmeConfig::from_env(MoeNvmeConfig base) { + base.host_slots = parse_bounded_int("DFLASH_MOE_NVME_SLOTS", base.host_slots, 2, 64); + base.io_threads = parse_bounded_int("DFLASH_MOE_NVME_IO_THREADS", base.io_threads, 1, 32); + base.demand_reserve = parse_bounded_int( + "DFLASH_MOE_NVME_DEMAND_RESERVE", base.demand_reserve, 1, base.host_slots - 1); + base.max_prefetch_batch = parse_bounded_int( + "DFLASH_MOE_NVME_PREFETCH_BATCH", base.max_prefetch_batch, 1, base.host_slots); + + if (const char * value = std::getenv("DFLASH_MOE_NVME_BACKEND")) { + const std::string mode = lowercase(value); + if (mode == "auto") base.backend = MoeNvmeBackend::Auto; + else if (mode == "thread" || mode == "threads" || mode == "pread") { + base.backend = MoeNvmeBackend::ThreadPool; + } else if (mode == "uring" || mode == "io_uring") { + base.backend = MoeNvmeBackend::IoUring; + } else if (mode == "mmap") { + base.backend = MoeNvmeBackend::Mmap; + } else { + std::fprintf(stderr, "[moe-nvme] ignoring invalid DFLASH_MOE_NVME_BACKEND=%s\n", value); + } + } + if (const char * value = std::getenv("DFLASH_MOE_NVME_DIRECT")) { + const std::string mode = lowercase(value); + if (mode == "auto") base.direct_io = MoeNvmeDirectMode::Auto; + else if (mode == "1" || mode == "on" || mode == "true") { + base.direct_io = MoeNvmeDirectMode::Enabled; + } else if (mode == "0" || mode == "off" || mode == "false") { + base.direct_io = MoeNvmeDirectMode::Disabled; + } else { + std::fprintf(stderr, "[moe-nvme] ignoring invalid DFLASH_MOE_NVME_DIRECT=%s\n", value); + } + } + return base; +} + +bool make_moe_expert_io_layout( + int layer, + int expert, + const LayerExpertRegions & regions, + size_t source_size, + size_t direct_alignment, + MoeExpertIoLayout & out, + std::string * err) { + + return make_moe_expert_io_layout( + layer, expert, regions, std::vector{source_size}, + direct_alignment, out, err); +} + +bool make_moe_expert_io_layout( + int layer, + int expert, + const LayerExpertRegions & regions, + const std::vector & source_sizes, + size_t direct_alignment, + MoeExpertIoLayout & out, + std::string * err) { + + out = {}; + out.key = { (int32_t) layer, (int32_t) expert }; + out.fused_gate_up = regions.fused_gate_up; + if (layer < 0 || expert < 0) { + if (err) *err = "negative layer or expert id"; + return false; + } + if (!is_power_of_two(direct_alignment) || direct_alignment < 512) { + if (err) *err = "direct I/O alignment must be a power of two >= 512"; + return false; + } + + size_t host_cursor = 0; + size_t device_cursor = 0; + auto add_span = [&](const ExpertFileRegion & region, size_t expert_bytes, + const char * label) -> bool { + if (out.span_count >= 3 || expert_bytes == 0 || region.size == 0) { + if (err) *err = std::string("missing or invalid ") + label + " expert region"; + return false; + } + if (region.source_index >= source_sizes.size()) { + if (err) *err = std::string(label) + " tensor references a missing model shard"; + return false; + } + const size_t source_size = source_sizes[region.source_index]; + if (!range_in_bounds(region.offset, region.size, source_size)) { + if (err) *err = std::string(label) + " tensor region is outside its model shard"; + return false; + } + const size_t expert_count = region.size / expert_bytes; + if ((size_t) expert >= expert_count) { + if (err) *err = std::string(label) + " expert id exceeds tensor extent"; + return false; + } + size_t expert_delta = 0; + size_t file_offset = 0; + if (!checked_mul((size_t) expert, expert_bytes, expert_delta) || + !checked_add(region.offset, expert_delta, file_offset) || + !range_in_bounds(expert_delta, expert_bytes, region.size) || + !range_in_bounds(file_offset, expert_bytes, source_size)) { + if (err) *err = std::string(label) + " expert range is outside the model file"; + return false; + } + + MoeExpertIoSpan & span = out.spans[out.span_count++]; + span.file_offset = file_offset; + span.source_index = region.source_index; + span.bytes = expert_bytes; + span.device_offset = device_cursor; + + const size_t aligned_file = align_down(file_offset, direct_alignment); + const size_t lead = file_offset - aligned_file; + size_t aligned_buffer = 0; + size_t raw_io_bytes = 0; + size_t aligned_io_bytes = 0; + if (!align_up_checked(host_cursor, direct_alignment, aligned_buffer) || + !checked_add(lead, expert_bytes, raw_io_bytes) || + !align_up_checked(raw_io_bytes, direct_alignment, aligned_io_bytes) || + !checked_add(aligned_buffer, lead, span.buffer_offset) || + !checked_add(aligned_buffer, aligned_io_bytes, host_cursor) || + !checked_add(device_cursor, expert_bytes, device_cursor)) { + if (err) *err = "expert I/O layout size overflow"; + return false; + } + span.io_file_offset = aligned_file; + span.io_buffer_offset = aligned_buffer; + span.io_bytes = aligned_io_bytes; + return true; + }; + + if (regions.fused_gate_up) { + if (!add_span(regions.gate_up_exps, regions.expert_bytes_gate_up, "gate_up")) return false; + } else { + if (!add_span(regions.gate_exps, regions.expert_bytes_gate, "gate")) return false; + if (!add_span(regions.up_exps, regions.expert_bytes_up, "up")) return false; + } + if (!add_span(regions.down_exps, regions.expert_bytes_down, "down")) return false; + + out.payload_bytes = device_cursor; + out.host_bytes = host_cursor; + return true; +} + +struct MoeNvmeScheduler::Impl { + enum class SlotState : uint8_t { Free, Queued, Reading, Ready, Failed }; + + struct Slot { + void * data = nullptr; + SlotState state = SlotState::Free; + MoeExpertKey key{}; + MoeExpertIoLayout layout{}; + MoeNvmePriority priority = MoeNvmePriority::Prefetch; + uint64_t generation = 0; + uint64_t queue_epoch = 0; + uint64_t last_touch = 0; + uint64_t frequency = 0; + int leases = 0; + bool demand_resident = false; + std::string error; + }; + + struct Job { + int slot = -1; + uint64_t generation = 0; + uint64_t queue_epoch = 0; + MoeNvmePriority priority = MoeNvmePriority::Prefetch; + }; + + struct SlotRef { + int slot = -1; + uint64_t generation = 0; + }; + + MoeNvmeConfig config{}; + size_t max_payload_bytes = 0; + size_t bytes_per_slot = 0; + AllocateFn allocate = nullptr; + FreeFn free_fn = nullptr; + void * allocator_opaque = nullptr; + std::vector slots; + std::vector regions; + std::vector sources; + std::vector source_sizes; + + MoeNvmeBackend effective_backend = MoeNvmeBackend::ThreadPool; + bool direct_active = false; + bool initialized = false; + bool bound = false; + bool stopping = false; + std::vector source_fds; + std::vector direct_fds; + std::vector active_fds; + + mutable std::mutex mutex; + std::condition_variable work_cv; + std::condition_variable state_cv; + std::deque demand_queue; + std::deque prefetch_queue; + std::unordered_map index; + std::vector workers; + uint64_t clock = 0; + int active_prefetch = 0; + int max_active_prefetch = 1; + +#if defined(__linux__) + std::unique_ptr ring; +#endif + + std::atomic requests{0}; + std::atomic demand_requests{0}; + std::atomic prefetch_requests{0}; + std::atomic cache_hits{0}; + std::atomic inflight_deduplications{0}; + std::atomic demand_upgrades{0}; + std::atomic prefetch_drops{0}; + std::atomic evictions{0}; + std::atomic read_ops{0}; + std::atomic payload_bytes{0}; + std::atomic physical_bytes{0}; + mutable std::mutex io_time_mutex; + int io_active = 0; + Clock::time_point io_active_begin{}; + uint64_t active_io_ns = 0; + std::atomic read_ns{0}; + std::atomic wait_ns{0}; + std::atomic errors{0}; + + void begin_io_activity() { + std::lock_guard lock(io_time_mutex); + if (io_active++ == 0) io_active_begin = Clock::now(); + } + + void end_io_activity() { + const auto now = Clock::now(); + std::lock_guard lock(io_time_mutex); + if (io_active <= 0) return; + if (--io_active == 0) active_io_ns += elapsed_ns(io_active_begin, now); + } + + uint64_t active_io_time() const { + const auto now = Clock::now(); + std::lock_guard lock(io_time_mutex); + return active_io_ns + (io_active > 0 ? elapsed_ns(io_active_begin, now) : 0); + } + + void reset_io_time() { + std::lock_guard lock(io_time_mutex); + active_io_ns = 0; + if (io_active > 0) io_active_begin = Clock::now(); + } + + bool valid_job_locked(const Job & job) const { + if (job.slot < 0 || job.slot >= (int) slots.size()) return false; + const Slot & slot = slots[(size_t) job.slot]; + return slot.state == SlotState::Queued && + slot.generation == job.generation && + slot.queue_epoch == job.queue_epoch; + } + + bool queue_has_valid_locked(std::deque & queue) { + while (!queue.empty() && !valid_job_locked(queue.front())) queue.pop_front(); + return !queue.empty(); + } + + bool take_one_locked(Job & out, bool allow_prefetch) { + if (queue_has_valid_locked(demand_queue)) { + out = demand_queue.front(); + demand_queue.pop_front(); + } else { + if (!allow_prefetch || active_prefetch >= max_active_prefetch || + !queue_has_valid_locked(prefetch_queue)) { + return false; + } + out = prefetch_queue.front(); + prefetch_queue.pop_front(); + ++active_prefetch; + } + Slot & slot = slots[(size_t) out.slot]; + slot.state = SlotState::Reading; + return true; + } + + int speculative_occupancy_locked() const { + int count = 0; + for (const Slot & slot : slots) { + if (slot.state != SlotState::Free && !slot.demand_resident && + slot.priority == MoeNvmePriority::Prefetch) { + ++count; + } + } + return count; + } + + uint64_t eviction_score_locked(const Slot & slot) const { + // Frequency dominates; recency breaks ties. This is an LFRU score with + // enough hysteresis that one recent speculative touch cannot displace + // a repeatedly demanded expert. + const uint64_t age = clock >= slot.last_touch ? clock - slot.last_touch : 0; + const uint64_t recency = age < 255 ? 255 - age : 0; + return (slot.frequency << 8) | recency; + } + + int choose_slot_locked(MoeNvmePriority priority) { + if (priority == MoeNvmePriority::Prefetch) { + const int limit = std::max(1, (int) slots.size() - config.demand_reserve); + if (speculative_occupancy_locked() >= limit) return -1; + } + for (size_t i = 0; i < slots.size(); ++i) { + if (slots[i].state == SlotState::Free) return (int) i; + } + + if (priority == MoeNvmePriority::Demand) { + // Demand may cancel a queued speculative read before evicting data. + for (size_t i = 0; i < slots.size(); ++i) { + Slot & slot = slots[i]; + if (slot.state == SlotState::Queued && + slot.priority == MoeNvmePriority::Prefetch && slot.leases == 0) { + index.erase(slot.key); + ++slot.queue_epoch; // invalidates the old queue entry + return (int) i; + } + } + } + + int victim = -1; + uint64_t best = std::numeric_limits::max(); + for (size_t i = 0; i < slots.size(); ++i) { + const Slot & slot = slots[i]; + if ((slot.state != SlotState::Ready && slot.state != SlotState::Failed) || + slot.leases != 0) { + continue; + } + if (priority == MoeNvmePriority::Prefetch && slot.demand_resident) continue; + uint64_t score = slot.state == SlotState::Failed ? 0 : eviction_score_locked(slot); + if (!slot.demand_resident) score >>= 2; + if (score < best) { + best = score; + victim = (int) i; + } + } + return victim; + } + + enum class Admission { New, ReadyHit, Inflight, NoSlot, Invalid }; + + Admission admit_locked(int layer, int expert, MoeNvmePriority priority, + int & slot_out, std::string * err) { + slot_out = -1; + if (!bound) { + if (err) *err = "SSD scheduler has no bound model source"; + return Admission::Invalid; + } + if (layer < 0 || layer >= (int) regions.size()) { + if (err) *err = "SSD request layer is out of range"; + return Admission::Invalid; + } + const MoeExpertKey key{ (int32_t) layer, (int32_t) expert }; + auto found = index.find(key); + if (found != index.end()) { + Slot & slot = slots[(size_t) found->second.slot]; + if (slot.generation == found->second.generation && slot.state != SlotState::Free) { + slot_out = found->second.slot; + slot.last_touch = ++clock; + if (slot.state == SlotState::Ready) { + ++slot.frequency; + if (priority == MoeNvmePriority::Demand) slot.demand_resident = true; + return Admission::ReadyHit; + } + if (slot.state == SlotState::Failed) { + if (err) *err = slot.error; + return Admission::Invalid; + } + if (priority == MoeNvmePriority::Demand && + slot.priority == MoeNvmePriority::Prefetch) { + slot.priority = MoeNvmePriority::Demand; + slot.demand_resident = true; + demand_upgrades.fetch_add(1, std::memory_order_relaxed); + if (slot.state == SlotState::Queued) { + ++slot.queue_epoch; + demand_queue.push_back({slot_out, slot.generation, + slot.queue_epoch, MoeNvmePriority::Demand}); + work_cv.notify_one(); + } + } + return Admission::Inflight; + } + index.erase(found); + } + + const int chosen = choose_slot_locked(priority); + if (chosen < 0) return Admission::NoSlot; + Slot & slot = slots[(size_t) chosen]; + if (slot.state != SlotState::Free) { + index.erase(slot.key); + evictions.fetch_add(1, std::memory_order_relaxed); + } + + MoeExpertIoLayout layout; + if (!make_moe_expert_io_layout(layer, expert, regions[(size_t) layer], + source_sizes, config.direct_alignment, + layout, err)) { + slot.state = SlotState::Free; + return Admission::Invalid; + } + if (layout.payload_bytes > max_payload_bytes || layout.host_bytes > bytes_per_slot) { + if (err) *err = "expert read plan exceeds the configured SSD slot size"; + slot.state = SlotState::Free; + return Admission::Invalid; + } + + slot.state = SlotState::Queued; + slot.key = key; + slot.layout = layout; + slot.priority = priority; + ++slot.generation; + ++slot.queue_epoch; + slot.last_touch = ++clock; + slot.frequency = priority == MoeNvmePriority::Demand ? 1 : 0; + slot.leases = 0; + slot.demand_resident = priority == MoeNvmePriority::Demand; + slot.error.clear(); + slot_out = chosen; + index[key] = {chosen, slot.generation}; + Job job{chosen, slot.generation, slot.queue_epoch, priority}; + if (priority == MoeNvmePriority::Demand) demand_queue.push_back(job); + else prefetch_queue.push_back(job); + work_cv.notify_one(); + return Admission::New; + } + + bool read_job_threaded(const Job & job, std::string & err, + uint64_t & ops, uint64_t & logical, uint64_t & physical) { + const Slot & slot = slots[(size_t) job.slot]; + uint8_t * base = static_cast(slot.data); + for (int i = 0; i < slot.layout.span_count; ++i) { + const MoeExpertIoSpan & span = slot.layout.spans[i]; + if (span.source_index >= sources.size()) { + err = "expert read references a missing model shard"; + return false; + } + const MoeNvmeSource & source = sources[span.source_index]; + const int active_fd = active_fds.empty() + ? -1 : active_fds[span.source_index]; + ++ops; + logical += span.bytes; + if (effective_backend == MoeNvmeBackend::Mmap || active_fd < 0) { + if (!source.mmap_data || + !range_in_bounds(span.file_offset, span.bytes, source.mmap_size)) { + err = "mmap expert read is outside the model file"; + return false; + } + const auto * src = static_cast(source.mmap_data); + std::memcpy(base + span.buffer_offset, src + span.file_offset, span.bytes); + physical += span.bytes; + } else { +#if defined(_WIN32) + (void) base; + err = "pread backend is unavailable on Windows"; + return false; +#else + const size_t read_offset = direct_active ? span.io_file_offset : span.file_offset; + const size_t read_bytes = direct_active ? span.io_bytes : span.bytes; + const size_t buffer_offset = direct_active ? span.io_buffer_offset : span.buffer_offset; + if (!pread_full(active_fd, base + buffer_offset, read_bytes, read_offset, err)) return false; + physical += read_bytes; +#endif + } + } + return true; + } + + void complete_job(const Job & job, bool ok, const std::string & err, + uint64_t ops, uint64_t logical, uint64_t physical, + uint64_t duration_ns) { + read_ops.fetch_add(ops, std::memory_order_relaxed); + payload_bytes.fetch_add(logical, std::memory_order_relaxed); + physical_bytes.fetch_add(physical, std::memory_order_relaxed); + read_ns.fetch_add(duration_ns, std::memory_order_relaxed); + if (!ok) errors.fetch_add(1, std::memory_order_relaxed); + + std::lock_guard lock(mutex); + if (job.priority == MoeNvmePriority::Prefetch && active_prefetch > 0) { + --active_prefetch; + } + if (job.slot >= 0 && job.slot < (int) slots.size()) { + Slot & slot = slots[(size_t) job.slot]; + if (slot.generation == job.generation && slot.state == SlotState::Reading) { + slot.state = ok ? SlotState::Ready : SlotState::Failed; + slot.error = ok ? std::string() : err; + slot.last_touch = ++clock; + } + } + state_cv.notify_all(); + work_cv.notify_all(); + } + + void thread_worker() { + for (;;) { + Job job; + { + std::unique_lock lock(mutex); + work_cv.wait(lock, [&] { + return stopping || queue_has_valid_locked(demand_queue) || + (active_prefetch < max_active_prefetch && + queue_has_valid_locked(prefetch_queue)); + }); + if (stopping) return; + if (!take_one_locked(job, true)) continue; + } + + const auto begin = Clock::now(); + std::string err; + uint64_t ops = 0, logical = 0, physical = 0; + begin_io_activity(); + const bool ok = read_job_threaded(job, err, ops, logical, physical); + end_io_activity(); + complete_job(job, ok, err, ops, logical, physical, + elapsed_ns(begin, Clock::now())); + } + } + +#if defined(__linux__) + void uring_worker() { + struct Op { + size_t job = 0; + uint32_t expected = 0; + }; + struct Progress { + int pending = 0; + uint64_t ops = 0; + uint64_t logical = 0; + uint64_t physical = 0; + bool ok = true; + bool completed = false; + std::string error; + }; + for (;;) { + std::vector jobs; + { + std::unique_lock lock(mutex); + work_cv.wait(lock, [&] { + return stopping || queue_has_valid_locked(demand_queue) || + queue_has_valid_locked(prefetch_queue); + }); + if (stopping) return; + + Job job; + while ((int) jobs.size() < config.host_slots && + take_one_locked(job, false)) { + jobs.push_back(job); + } + if (jobs.empty()) { + int speculative = 0; + while ((int) jobs.size() < config.max_prefetch_batch && + speculative < config.max_prefetch_batch && + take_one_locked(job, true)) { + jobs.push_back(job); + ++speculative; + } + } + } + if (jobs.empty()) continue; + + const auto begin = Clock::now(); + std::vector operations; + std::vector progress(jobs.size()); + + for (size_t j = 0; j < jobs.size(); ++j) { + const Slot & slot = slots[(size_t) jobs[j].slot]; + auto * base = static_cast(slot.data); + progress[j].logical = slot.layout.payload_bytes; + for (int s = 0; s < slot.layout.span_count; ++s) { + const MoeExpertIoSpan & span = slot.layout.spans[s]; + const size_t read_offset = direct_active ? span.io_file_offset : span.file_offset; + const size_t read_bytes = direct_active ? span.io_bytes : span.bytes; + const size_t buffer_offset = direct_active ? span.io_buffer_offset : span.buffer_offset; + if (read_bytes > std::numeric_limits::max()) { + progress[j].ok = false; + progress[j].error = + "one expert tensor read exceeds io_uring's 32-bit length"; + continue; + } + io_uring_sqe * sqe = ring->get_sqe(); + if (!sqe) { + progress[j].ok = false; + progress[j].error = "io_uring submission queue is full"; + continue; + } + const uint64_t op_index = operations.size(); + operations.push_back({j, (uint32_t) read_bytes}); + ring->prepare_read(sqe, active_fds, span.source_index, + jobs[j].slot, + base + buffer_offset, (uint32_t) read_bytes, + read_offset, op_index, !direct_active); + ++progress[j].pending; + ++progress[j].ops; + progress[j].physical += read_bytes; + } + } + + // Publish a slot as soon as that expert's last tensor slice + // completes. Waiting for the slowest expert in the whole ring + // batch creates a barrier that prevents SSD N+1 from overlapping + // H2D/compute N. + auto finish_job = [&](size_t j) { + Progress & item = progress[j]; + if (item.completed) return; + item.completed = true; + complete_job(jobs[j], item.ok, item.error, item.ops, + item.logical, item.physical, + elapsed_ns(begin, Clock::now())); + }; + + std::string ring_error; + bool fatal_ring_error = false; + if (!operations.empty()) begin_io_activity(); + if (operations.empty()) { + for (Progress & item : progress) { + item.ok = false; + if (item.error.empty()) item.error = "io_uring batch contained no readable spans"; + } + } else if (!ring->submit_all(ring_error)) { + fatal_ring_error = true; + for (Progress & item : progress) { + item.ok = false; + item.error = ring_error; + } + } else { + for (size_t j = 0; j < progress.size(); ++j) { + if (progress[j].pending == 0) finish_job(j); + } + for (size_t completed = 0; completed < operations.size(); ++completed) { + io_uring_cqe cqe{}; + if (!ring->wait_cqe(cqe, ring_error)) { + fatal_ring_error = true; + for (Progress & item : progress) { + if (!item.completed) { + item.ok = false; + item.error = ring_error; + } + } + break; + } + if (cqe.user_data >= operations.size()) { + for (Progress & item : progress) { + if (!item.completed) { + item.ok = false; + item.error = "io_uring returned an invalid completion tag"; + } + } + continue; + } + const Op & op = operations[(size_t) cqe.user_data]; + Progress & item = progress[op.job]; + if (cqe.res != (int32_t) op.expected) { + item.ok = false; + if (cqe.res < 0) { + item.error = std::string("io_uring read failed: ") + + std::strerror(-cqe.res); + } else { + item.error = "io_uring returned a short model read"; + } + } + if (item.pending > 0 && --item.pending == 0) finish_job(op.job); + } + } + if (!operations.empty()) end_io_activity(); + + for (size_t j = 0; j < jobs.size(); ++j) { + if (!progress[j].completed) { + if (progress[j].pending != 0 && progress[j].error.empty()) { + progress[j].ok = false; + progress[j].error = "io_uring did not complete every expert span"; + } + finish_job(j); + } + } + if (fatal_ring_error) { + std::lock_guard lock(mutex); + stopping = true; + state_cv.notify_all(); + work_cv.notify_all(); + return; + } + } + } +#endif +}; + +MoeNvmeLease::~MoeNvmeLease() { reset(); } + +MoeNvmeLease::MoeNvmeLease(MoeNvmeLease && other) noexcept + : scheduler_(other.scheduler_), data_(other.data_), layout_(other.layout_), + slot_(other.slot_), generation_(other.generation_) { + other.scheduler_ = nullptr; + other.data_ = nullptr; + other.slot_ = -1; + other.generation_ = 0; +} + +MoeNvmeLease & MoeNvmeLease::operator=(MoeNvmeLease && other) noexcept { + if (this != &other) { + reset(); + scheduler_ = other.scheduler_; + data_ = other.data_; + layout_ = other.layout_; + slot_ = other.slot_; + generation_ = other.generation_; + other.scheduler_ = nullptr; + other.data_ = nullptr; + other.slot_ = -1; + other.generation_ = 0; + } + return *this; +} + +void MoeNvmeLease::reset() { + if (scheduler_) scheduler_->release_lease(slot_, generation_); + scheduler_ = nullptr; + data_ = nullptr; + slot_ = -1; + generation_ = 0; + layout_ = {}; +} + +MoeNvmeScheduler::MoeNvmeScheduler() : impl_(new Impl) {} +MoeNvmeScheduler::~MoeNvmeScheduler() { destroy(); } + +bool MoeNvmeScheduler::init(const MoeNvmeConfig & requested, + size_t max_expert_payload_bytes, + AllocateFn allocate, + FreeFn free_fn, + void * allocator_opaque, + std::string * err) { + destroy(); + impl_.reset(new (std::nothrow) Impl); + if (!impl_) { + if (err) *err = "failed to allocate SSD scheduler state"; + return false; + } + Impl & p = *impl_; + p.config = requested; + p.config.host_slots = std::max(2, p.config.host_slots); + p.config.io_threads = std::max(1, p.config.io_threads); + p.config.demand_reserve = std::max(1, std::min(p.config.demand_reserve, + p.config.host_slots - 1)); + p.config.max_prefetch_batch = std::max(1, std::min(p.config.max_prefetch_batch, + p.config.host_slots)); + if (!is_power_of_two(p.config.direct_alignment) || p.config.direct_alignment < 512 || + max_expert_payload_bytes == 0 || !allocate || !free_fn) { + if (err) *err = "invalid SSD scheduler initialization arguments"; + return false; + } + p.max_payload_bytes = max_expert_payload_bytes; + p.allocate = allocate; + p.free_fn = free_fn; + p.allocator_opaque = allocator_opaque; + + size_t overhead = 0; + if (!checked_mul((size_t) 16, p.config.direct_alignment, overhead) || + !checked_add(max_expert_payload_bytes, overhead, p.bytes_per_slot) || + !align_up_checked(p.bytes_per_slot, p.config.direct_alignment, p.bytes_per_slot)) { + if (err) *err = "SSD slot size overflow"; + return false; + } + + p.slots.resize((size_t) p.config.host_slots); + for (size_t i = 0; i < p.slots.size(); ++i) { + if (!p.allocate(&p.slots[i].data, p.bytes_per_slot, p.allocator_opaque) || + !p.slots[i].data) { + if (err) *err = "failed to allocate page-locked SSD host slot"; + for (size_t j = 0; j < i; ++j) p.free_fn(p.slots[j].data, p.allocator_opaque); + p.slots.clear(); + return false; + } + } + p.initialized = true; + return true; +} + +bool MoeNvmeScheduler::bind_source(const MoeNvmeSource & source, + const std::vector & layer_regions, + std::string * err) { + return bind_sources({source}, layer_regions, err); +} + +bool MoeNvmeScheduler::bind_sources( + const std::vector & sources, + const std::vector & layer_regions, + std::string * err) { + if (!impl_ || !impl_->initialized) { + if (err) *err = "SSD scheduler is not initialized"; + return false; + } + Impl & p = *impl_; + std::lock_guard lock(p.mutex); + if (p.bound) { + bool same = p.sources.size() == sources.size(); + for (size_t i = 0; same && i < sources.size(); ++i) { + same = p.sources[i].mmap_data == sources[i].mmap_data && + p.sources[i].mmap_size == sources[i].mmap_size && + (sources[i].mmap_data || p.sources[i].fd == sources[i].fd); + } + if (!same && err) *err = "SSD scheduler cannot rebind an active model source"; + return same; + } + if (sources.empty() || layer_regions.empty()) { + if (err) *err = "empty SSD model source or expert-region table"; + return false; + } + bool all_mapped = true; + bool all_have_fds = true; + uint64_t total_source_bytes = 0; + for (const MoeNvmeSource & source : sources) { + if (source.mmap_size == 0 || (!source.mmap_data && source.fd < 0)) { + if (err) *err = "invalid SSD model shard"; + return false; + } + all_mapped = all_mapped && source.mmap_data != nullptr; + all_have_fds = all_have_fds && source.fd >= 0; + if ((uint64_t) source.mmap_size > + std::numeric_limits::max() - total_source_bytes) { + if (err) *err = "SSD model shard sizes overflow"; + return false; + } + total_source_bytes += (uint64_t) source.mmap_size; + } + p.sources = sources; + p.source_sizes.clear(); + p.source_sizes.reserve(sources.size()); + for (const MoeNvmeSource & source : sources) { + p.source_sizes.push_back(source.mmap_size); + } + p.regions = layer_regions; + p.source_fds.assign(sources.size(), -1); + p.direct_fds.assign(sources.size(), -1); + + auto close_model_fds = [&]() { +#if !defined(_WIN32) + for (int fd : p.direct_fds) if (fd >= 0) ::close(fd); + for (int fd : p.source_fds) if (fd >= 0) ::close(fd); +#endif + p.direct_fds.clear(); + p.source_fds.clear(); + p.active_fds.clear(); + p.direct_active = false; + }; + +#if !defined(_WIN32) + if (all_have_fds) { + for (size_t i = 0; i < sources.size(); ++i) { + p.source_fds[i] = ::dup(sources[i].fd); + if (p.source_fds[i] < 0) { + if (err) { + *err = std::string("failed to duplicate model shard fd: ") + + std::strerror(errno); + } + close_model_fds(); + return false; + } + } + } +#endif + const bool duplicated_all_fds = all_have_fds && + std::all_of(p.source_fds.begin(), p.source_fds.end(), + [](int fd) { return fd >= 0; }); + + bool want_direct = p.config.direct_io == MoeNvmeDirectMode::Enabled; + if (p.config.direct_io == MoeNvmeDirectMode::Auto) { + const uint64_t ram = physical_memory_bytes(); + // Direct I/O avoids keeping both an explicit expert cache and a second + // model-sized page cache when the model itself nearly fills RAM. + want_direct = ram != 0 && total_source_bytes > ram - ram / 4; + } + if (want_direct && !duplicated_all_fds && + p.config.direct_io == MoeNvmeDirectMode::Enabled) { + if (err) *err = "O_DIRECT requires a readable fd for every model shard"; + close_model_fds(); + return false; + } + +#if defined(__linux__) && defined(O_DIRECT) + if (want_direct && duplicated_all_fds) { + bool aligned = true; + for (const Impl::Slot & slot : p.slots) { + if (((uintptr_t) slot.data & (p.config.direct_alignment - 1)) != 0) { + aligned = false; + break; + } + } + if (aligned) { + p.direct_active = true; + for (size_t i = 0; i < p.source_fds.size(); ++i) { + char proc_path[64]; + std::snprintf(proc_path, sizeof(proc_path), + "/proc/self/fd/%d", p.source_fds[i]); + p.direct_fds[i] = ::open( + proc_path, O_RDONLY | O_CLOEXEC | O_DIRECT); + if (p.direct_fds[i] < 0) { + p.direct_active = false; + break; + } + } + if (!p.direct_active) { + for (int & fd : p.direct_fds) { + if (fd >= 0) ::close(fd); + fd = -1; + } + } + } + if (!p.direct_active && p.config.direct_io == MoeNvmeDirectMode::Enabled) { + if (err) { + *err = "O_DIRECT was requested but one model shard or the pinned buffers do not support it"; + } + close_model_fds(); + return false; + } + } +#else + if (want_direct && p.config.direct_io == MoeNvmeDirectMode::Enabled) { + if (err) *err = "O_DIRECT was requested on an unsupported platform"; + close_model_fds(); + return false; + } +#endif + if (duplicated_all_fds) { + p.active_fds = p.direct_active ? p.direct_fds : p.source_fds; + } + + if (p.config.backend == MoeNvmeBackend::Mmap) { + if (!all_mapped) { + if (err) *err = "mmap SSD backend requested without every shard mapped"; + close_model_fds(); + return false; + } + p.effective_backend = MoeNvmeBackend::Mmap; + close_model_fds(); + } else if (p.config.backend == MoeNvmeBackend::ThreadPool) { + if (!p.active_fds.empty()) p.effective_backend = MoeNvmeBackend::ThreadPool; + else if (all_mapped) p.effective_backend = MoeNvmeBackend::Mmap; + else { + if (err) *err = "threaded SSD backend needs every shard readable"; + close_model_fds(); + return false; + } + } else { +#if defined(__linux__) + if (!p.active_fds.empty()) { + p.ring.reset(new (std::nothrow) RawIoUring); + std::string ring_error; + std::vector buffers; + buffers.reserve(p.slots.size()); + for (const Impl::Slot & slot : p.slots) buffers.push_back(slot.data); + const unsigned entries = (unsigned) std::max( + 32, p.slots.size() * 3 + 4); + if (p.ring && p.ring->open(entries, p.active_fds, buffers, + p.bytes_per_slot, ring_error)) { + p.effective_backend = MoeNvmeBackend::IoUring; + } else { + p.ring.reset(); + if (p.config.backend == MoeNvmeBackend::IoUring) { + if (err) *err = ring_error.empty() ? "io_uring initialization failed" : ring_error; + close_model_fds(); + return false; + } + p.effective_backend = MoeNvmeBackend::ThreadPool; + } + } else { + if (p.config.backend == MoeNvmeBackend::IoUring) { + if (err) *err = "io_uring backend requested without every model shard readable"; + close_model_fds(); + return false; + } + if (all_mapped) p.effective_backend = MoeNvmeBackend::Mmap; + else { + if (err) *err = "SSD backend cannot read every model shard"; + close_model_fds(); + return false; + } + } +#else + if (p.config.backend == MoeNvmeBackend::IoUring) { + if (err) *err = "io_uring backend is Linux-only"; + close_model_fds(); + return false; + } + if (!p.active_fds.empty()) p.effective_backend = MoeNvmeBackend::ThreadPool; + else if (all_mapped) p.effective_backend = MoeNvmeBackend::Mmap; + else { + if (err) *err = "SSD backend cannot read every model shard"; + close_model_fds(); + return false; + } +#endif + } + + p.stopping = false; + if (p.effective_backend == MoeNvmeBackend::IoUring) { +#if defined(__linux__) + p.max_active_prefetch = std::max(1, p.config.max_prefetch_batch); + p.workers.emplace_back([&p] { p.uring_worker(); }); +#endif + } else { + const int workers = p.effective_backend == MoeNvmeBackend::Mmap + ? std::min(2, p.config.io_threads) : p.config.io_threads; + p.max_active_prefetch = std::max(1, workers - 1); + for (int i = 0; i < workers; ++i) { + p.workers.emplace_back([&p] { p.thread_worker(); }); + } + } + p.bound = true; + return true; +} + +bool MoeNvmeScheduler::is_initialized() const { + return impl_ && impl_->initialized; +} + +bool MoeNvmeScheduler::is_bound() const { + return impl_ && impl_->bound; +} + +void MoeNvmeScheduler::destroy() { + if (!impl_) return; + Impl & p = *impl_; + { + std::lock_guard lock(p.mutex); + p.stopping = true; + p.work_cv.notify_all(); + p.state_cv.notify_all(); + } + for (std::thread & worker : p.workers) { + if (worker.joinable()) worker.join(); + } + p.workers.clear(); +#if defined(__linux__) + p.ring.reset(); +#endif +#if !defined(_WIN32) + for (int fd : p.direct_fds) if (fd >= 0) ::close(fd); + for (int fd : p.source_fds) if (fd >= 0) ::close(fd); +#endif + p.direct_fds.clear(); + p.source_fds.clear(); + p.active_fds.clear(); + for (Impl::Slot & slot : p.slots) { + if (slot.data && p.free_fn) p.free_fn(slot.data, p.allocator_opaque); + slot.data = nullptr; + } + p.slots.clear(); + p.index.clear(); + p.demand_queue.clear(); + p.prefetch_queue.clear(); + p.regions.clear(); + p.sources.clear(); + p.source_sizes.clear(); + p.initialized = false; + p.bound = false; +} + +bool MoeNvmeScheduler::request(int layer, int expert, MoeNvmePriority priority, + std::string * err) { + if (!impl_ || !impl_->initialized) { + if (err) *err = "SSD scheduler is not initialized"; + return false; + } + Impl & p = *impl_; + p.requests.fetch_add(1, std::memory_order_relaxed); + if (priority == MoeNvmePriority::Demand) { + p.demand_requests.fetch_add(1, std::memory_order_relaxed); + } else { + p.prefetch_requests.fetch_add(1, std::memory_order_relaxed); + } + std::lock_guard lock(p.mutex); + int slot = -1; + const Impl::Admission result = p.admit_locked(layer, expert, priority, slot, err); + if (result == Impl::Admission::ReadyHit) { + p.cache_hits.fetch_add(1, std::memory_order_relaxed); + return true; + } + if (result == Impl::Admission::Inflight) { + p.inflight_deduplications.fetch_add(1, std::memory_order_relaxed); + return true; + } + if (result == Impl::Admission::New) return true; + if (result == Impl::Admission::NoSlot && priority == MoeNvmePriority::Prefetch) { + p.prefetch_drops.fetch_add(1, std::memory_order_relaxed); + return false; + } + if (result == Impl::Admission::NoSlot && err) *err = "all SSD expert slots are busy"; + return false; +} + +bool MoeNvmeScheduler::acquire(int layer, int expert, MoeNvmeLease & out, + std::string * err) { + out.reset(); + if (!impl_ || !impl_->initialized) { + if (err) *err = "SSD scheduler is not initialized"; + return false; + } + Impl & p = *impl_; + const auto begin = Clock::now(); + p.requests.fetch_add(1, std::memory_order_relaxed); + p.demand_requests.fetch_add(1, std::memory_order_relaxed); + const MoeExpertKey key{(int32_t) layer, (int32_t) expert}; + + std::unique_lock lock(p.mutex); + bool admitted = false; + for (;;) { + if (p.stopping) { + if (err) *err = "SSD scheduler is stopping"; + return false; + } + auto found = p.index.find(key); + if (found == p.index.end()) { + int slot = -1; + const Impl::Admission result = p.admit_locked( + layer, expert, MoeNvmePriority::Demand, slot, err); + if (result == Impl::Admission::Invalid) return false; + if (result == Impl::Admission::NoSlot) { + p.state_cv.wait(lock); + continue; + } + admitted = true; + if (result == Impl::Admission::ReadyHit) { + p.cache_hits.fetch_add(1, std::memory_order_relaxed); + } else if (result == Impl::Admission::Inflight) { + p.inflight_deduplications.fetch_add(1, std::memory_order_relaxed); + } + found = p.index.find(key); + if (found == p.index.end()) continue; + } + + Impl::Slot & slot = p.slots[(size_t) found->second.slot]; + if (slot.generation != found->second.generation) { + p.index.erase(found); + continue; + } + if (!admitted) { + if (slot.state == Impl::SlotState::Ready) { + p.cache_hits.fetch_add(1, std::memory_order_relaxed); + } else if (slot.state == Impl::SlotState::Queued || + slot.state == Impl::SlotState::Reading) { + p.inflight_deduplications.fetch_add(1, std::memory_order_relaxed); + } + admitted = true; + } + if (slot.priority == MoeNvmePriority::Prefetch && + (slot.state == Impl::SlotState::Queued || slot.state == Impl::SlotState::Reading)) { + slot.priority = MoeNvmePriority::Demand; + slot.demand_resident = true; + p.demand_upgrades.fetch_add(1, std::memory_order_relaxed); + if (slot.state == Impl::SlotState::Queued) { + ++slot.queue_epoch; + p.demand_queue.push_back({found->second.slot, slot.generation, + slot.queue_epoch, MoeNvmePriority::Demand}); + p.work_cv.notify_one(); + } + } + if (slot.state == Impl::SlotState::Failed) { + if (err) *err = slot.error; + return false; + } + if (slot.state != Impl::SlotState::Ready) { + p.state_cv.wait(lock); + continue; + } + + ++slot.leases; + ++slot.frequency; + slot.last_touch = ++p.clock; + slot.demand_resident = true; + out.scheduler_ = this; + out.data_ = static_cast(slot.data); + out.layout_ = slot.layout; + out.slot_ = found->second.slot; + out.generation_ = slot.generation; + p.wait_ns.fetch_add(elapsed_ns(begin, Clock::now()), std::memory_order_relaxed); + return true; + } +} + +void MoeNvmeScheduler::release_lease(int slot_index, uint64_t generation) { + if (!impl_) return; + Impl & p = *impl_; + std::lock_guard lock(p.mutex); + if (slot_index >= 0 && slot_index < (int) p.slots.size()) { + Impl::Slot & slot = p.slots[(size_t) slot_index]; + if (slot.generation == generation && slot.leases > 0) --slot.leases; + } + p.state_cv.notify_all(); +} + +MoeNvmeStats MoeNvmeScheduler::stats() const { + MoeNvmeStats out; + if (!impl_) return out; + const Impl & p = *impl_; + out.requests = p.requests.load(std::memory_order_relaxed); + out.demand_requests = p.demand_requests.load(std::memory_order_relaxed); + out.prefetch_requests = p.prefetch_requests.load(std::memory_order_relaxed); + out.cache_hits = p.cache_hits.load(std::memory_order_relaxed); + out.inflight_deduplications = p.inflight_deduplications.load(std::memory_order_relaxed); + out.demand_upgrades = p.demand_upgrades.load(std::memory_order_relaxed); + out.prefetch_drops = p.prefetch_drops.load(std::memory_order_relaxed); + out.evictions = p.evictions.load(std::memory_order_relaxed); + out.read_ops = p.read_ops.load(std::memory_order_relaxed); + out.payload_bytes = p.payload_bytes.load(std::memory_order_relaxed); + out.physical_bytes = p.physical_bytes.load(std::memory_order_relaxed); + out.active_io_ns = p.active_io_time(); + out.read_ns = p.read_ns.load(std::memory_order_relaxed); + out.wait_ns = p.wait_ns.load(std::memory_order_relaxed); + out.errors = p.errors.load(std::memory_order_relaxed); + return out; +} + +void MoeNvmeScheduler::reset_stats() { + if (!impl_) return; + Impl & p = *impl_; + p.requests.store(0, std::memory_order_relaxed); + p.demand_requests.store(0, std::memory_order_relaxed); + p.prefetch_requests.store(0, std::memory_order_relaxed); + p.cache_hits.store(0, std::memory_order_relaxed); + p.inflight_deduplications.store(0, std::memory_order_relaxed); + p.demand_upgrades.store(0, std::memory_order_relaxed); + p.prefetch_drops.store(0, std::memory_order_relaxed); + p.evictions.store(0, std::memory_order_relaxed); + p.read_ops.store(0, std::memory_order_relaxed); + p.payload_bytes.store(0, std::memory_order_relaxed); + p.physical_bytes.store(0, std::memory_order_relaxed); + p.reset_io_time(); + p.read_ns.store(0, std::memory_order_relaxed); + p.wait_ns.store(0, std::memory_order_relaxed); + p.errors.store(0, std::memory_order_relaxed); +} + +size_t MoeNvmeScheduler::slot_bytes() const { + return impl_ ? impl_->bytes_per_slot : 0; +} + +size_t MoeNvmeScheduler::total_host_bytes() const { + return impl_ ? impl_->bytes_per_slot * impl_->slots.size() : 0; +} + +int MoeNvmeScheduler::slot_count() const { + return impl_ ? (int) impl_->slots.size() : 0; +} + +const char * MoeNvmeScheduler::effective_backend_name() const { + if (!impl_ || !impl_->bound) return "unbound"; + switch (impl_->effective_backend) { + case MoeNvmeBackend::IoUring: return impl_->direct_active ? "io_uring+direct" : "io_uring"; + case MoeNvmeBackend::ThreadPool: return impl_->direct_active ? "pread-pool+direct" : "pread-pool"; + case MoeNvmeBackend::Mmap: return "mmap-workers"; + case MoeNvmeBackend::Auto: break; + } + return "unknown"; +} + +bool MoeNvmeScheduler::direct_io_active() const { + return impl_ && impl_->direct_active; +} + +} // namespace dflash::common diff --git a/server/src/common/moe_nvme_scheduler.h b/server/src/common/moe_nvme_scheduler.h new file mode 100644 index 000000000..4e3537a36 --- /dev/null +++ b/server/src/common/moe_nvme_scheduler.h @@ -0,0 +1,231 @@ +// Model-neutral asynchronous SSD scheduler for routed MoE weights. +// +// This layer deliberately knows nothing about a model architecture or ggml. +// Callers describe an expert as two or three byte ranges in a model file. The +// scheduler owns a bounded set of page-locked host slots, merges duplicate +// requests, gives demand reads strict priority over speculation, and retains +// completed demand reads as a small protected cache. + +#pragma once + +#include "moe_hybrid_storage.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +enum class MoeNvmeBackend { + Auto, + ThreadPool, + IoUring, + Mmap, +}; + +enum class MoeNvmeDirectMode { + Auto, + Disabled, + Enabled, +}; + +enum class MoeNvmePriority : uint8_t { + Prefetch = 0, + Demand = 1, +}; + +struct MoeNvmeConfig { + // Host slots are both the in-flight queue bound and the small L2 cache. + // Eight slots give one modern NVMe enough queue depth without consuming a + // model-sized amount of pinned memory. + int host_slots = 8; + int io_threads = 4; + + MoeNvmeBackend backend = MoeNvmeBackend::Auto; + MoeNvmeDirectMode direct_io = MoeNvmeDirectMode::Auto; + + // At least this many slots cannot be occupied by speculative requests. + int demand_reserve = 2; + + // Limit speculative work admitted in one io_uring batch. This bounds the + // latency of a demand arriving immediately after speculation was issued. + int max_prefetch_batch = 2; + + size_t direct_alignment = 4096; + + // Environment overrides use the DFLASH_MOE_NVME_* prefix. Invalid values + // leave the supplied/default value unchanged. + static MoeNvmeConfig from_env(); + static MoeNvmeConfig from_env(MoeNvmeConfig base); +}; + +struct MoeNvmeSource { + const void * mmap_data = nullptr; + size_t mmap_size = 0; + int fd = -1; // borrowed; bind_source(s) duplicates it on POSIX +}; + +struct MoeExpertKey { + int32_t layer = -1; + int32_t expert = -1; + + bool operator==(const MoeExpertKey & other) const { + return layer == other.layer && expert == other.expert; + } +}; + +// One logical tensor slice and its direct-I/O envelope. +struct MoeExpertIoSpan { + size_t file_offset = 0; // first payload byte in the model file + uint32_t source_index = 0; // model shard containing this span + size_t bytes = 0; // payload bytes + size_t buffer_offset = 0; // first payload byte in the host slot + size_t device_offset = 0; // packed destination offset on the GPU + + size_t io_file_offset = 0; // aligned direct-I/O start + size_t io_buffer_offset = 0; // aligned direct-I/O destination + size_t io_bytes = 0; // aligned direct-I/O length +}; + +struct MoeExpertIoLayout { + MoeExpertKey key; + MoeExpertIoSpan spans[3]{}; + int span_count = 0; + size_t payload_bytes = 0; + size_t host_bytes = 0; + bool fused_gate_up = false; +}; + +// Convert the common LayerExpertRegions descriptor into an exact read plan. +// The plan is usable by mmap, buffered pread, and aligned direct I/O. +bool make_moe_expert_io_layout( + int layer, + int expert, + const LayerExpertRegions & regions, + size_t source_size, + size_t direct_alignment, + MoeExpertIoLayout & out, + std::string * err = nullptr); + +// Split-model variant. Each ExpertFileRegion selects one entry by +// source_index; a tensor itself is never split across files. +bool make_moe_expert_io_layout( + int layer, + int expert, + const LayerExpertRegions & regions, + const std::vector & source_sizes, + size_t direct_alignment, + MoeExpertIoLayout & out, + std::string * err = nullptr); + +struct MoeNvmeStats { + uint64_t requests = 0; + uint64_t demand_requests = 0; + uint64_t prefetch_requests = 0; + uint64_t cache_hits = 0; + uint64_t inflight_deduplications = 0; + uint64_t demand_upgrades = 0; + uint64_t prefetch_drops = 0; + uint64_t evictions = 0; + uint64_t read_ops = 0; + uint64_t payload_bytes = 0; + uint64_t physical_bytes = 0; + // Union of intervals in which at least one storage operation was active. + // Unlike read_ns, this does not double-count concurrent reads. + uint64_t active_io_ns = 0; + uint64_t read_ns = 0; + uint64_t wait_ns = 0; + uint64_t errors = 0; +}; + +class MoeNvmeScheduler; + +// A lease pins one host slot against eviction while an H2D transfer consumes +// it. It is intentionally move-only. +class MoeNvmeLease { +public: + MoeNvmeLease() = default; + ~MoeNvmeLease(); + + MoeNvmeLease(const MoeNvmeLease &) = delete; + MoeNvmeLease & operator=(const MoeNvmeLease &) = delete; + MoeNvmeLease(MoeNvmeLease && other) noexcept; + MoeNvmeLease & operator=(MoeNvmeLease && other) noexcept; + + explicit operator bool() const { return scheduler_ != nullptr; } + const uint8_t * data() const { return data_; } + const MoeExpertIoLayout & layout() const { return layout_; } + int slot_index() const { return slot_; } + void reset(); + +private: + friend class MoeNvmeScheduler; + MoeNvmeScheduler * scheduler_ = nullptr; + const uint8_t * data_ = nullptr; + MoeExpertIoLayout layout_{}; + int slot_ = -1; + uint64_t generation_ = 0; +}; + +class MoeNvmeScheduler { +public: + using AllocateFn = bool (*)(void ** ptr, size_t bytes, void * opaque); + using FreeFn = void (*)(void * ptr, void * opaque); + + MoeNvmeScheduler(); + ~MoeNvmeScheduler(); + + MoeNvmeScheduler(const MoeNvmeScheduler &) = delete; + MoeNvmeScheduler & operator=(const MoeNvmeScheduler &) = delete; + + // Allocate the bounded slot pool. bind_source() starts I/O workers after + // the model file and per-layer layouts are available. + bool init(const MoeNvmeConfig & config, + size_t max_expert_payload_bytes, + AllocateFn allocate, + FreeFn free_fn, + void * allocator_opaque, + std::string * err = nullptr); + + bool bind_source(const MoeNvmeSource & source, + const std::vector & layer_regions, + std::string * err = nullptr); + + bool bind_sources(const std::vector & sources, + const std::vector & layer_regions, + std::string * err = nullptr); + + bool is_initialized() const; + bool is_bound() const; + void destroy(); + + // Non-blocking admission. A false prefetch result simply means the bounded + // speculative budget was full; a false demand result includes an error. + bool request(int layer, int expert, MoeNvmePriority priority, + std::string * err = nullptr); + + // Request (or upgrade) and wait for an exact expert. On success the lease + // protects the host slot until reset/destruction. + bool acquire(int layer, int expert, MoeNvmeLease & out, + std::string * err = nullptr); + + MoeNvmeStats stats() const; + void reset_stats(); + + size_t slot_bytes() const; + size_t total_host_bytes() const; + int slot_count() const; + const char * effective_backend_name() const; + bool direct_io_active() const; + +private: + friend class MoeNvmeLease; + void release_lease(int slot, uint64_t generation); + + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index e6702903b..e530a438e 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -728,16 +728,21 @@ bool DeepSeek4Backend::init_moe_tensor_parallel() { } if (ds4_inprocess_moe_tp_enabled()) { - if (!expert_backend_ || !moe_hybrid_->materialized_cold_experts || - moe_hybrid_->cold_backend != expert_backend_) { + const bool resident_ready = moe_hybrid_->materialized_cold_experts && + moe_hybrid_->cold_backend == expert_backend_; + const bool streamed_ready = !moe_hybrid_->materialized_cold_experts && + stream_engine_.is_bound() && + stream_engine_.compute_backend() == expert_backend_; + if (!expert_backend_ || (!resident_ready && !streamed_ready)) { std::fprintf(stderr, "[deepseek4-moe-tp] in-process expert backend is not ready\n"); return false; } expert_runtime_.reset(); std::fprintf(stderr, - "[deepseek4-moe-tp] enabled mode=in-process local_gpu=%d " + "[deepseek4-moe-tp] enabled mode=%s local_gpu=%d " "expert_gpu=%d local_experts=%d remote_experts=%d\n", + streamed_ready ? "in-process+ssd" : "in-process", cfg_.device.gpu, ds4_moe_tp_gpu(cfg_.device.gpu), moe_placement_.total_hot, w_.n_layer * w_.n_expert - moe_placement_.total_hot); @@ -853,6 +858,7 @@ bool DeepSeek4Backend::init_hybrid_model() { auto hybrid = std::make_shared(); MoeHybridConfig hybrid_cfg = make_ds4_parent_worker_cfg(w_); + size_t nvme_device_cache_bytes = 0; const bool inprocess_tp = env_flag_enabled("DFLASH_DS4_MOE_TP") && ds4_inprocess_moe_tp_enabled(); if (inprocess_tp) { @@ -875,8 +881,47 @@ bool DeepSeek4Backend::init_hybrid_model() { expert_gpu); return false; } - hybrid_cfg.materialize_cold_experts = true; hybrid_cfg.cold_expert_backend = MoeHybridColdBackend::Gpu; + + // Third tier: R9700 static-hot -> Strix adaptive-warm -> SSD misses. + // Keep the established fully-resident path whenever it fits. In auto + // mode, switch only when the cold stack plus a safety reserve exceeds + // currently available Strix memory. An explicit override is useful + // for qualification on models that already fit. + bool stream_cold = false; + const char * nvme_mode = std::getenv("DFLASH_MOE_NVME_COLD_TIER"); + if (nvme_mode && (std::strcmp(nvme_mode, "1") == 0 || + std::strcmp(nvme_mode, "on") == 0 || + std::strcmp(nvme_mode, "true") == 0)) { + stream_cold = true; + } else if (!nvme_mode || std::strcmp(nvme_mode, "auto") == 0) { + Ds4ExpertMemoryInfo info; + std::string memory_error; + size_t expert_free = 0; + size_t expert_total = 0; + ggml_backend_cuda_get_device_memory(expert_gpu, &expert_free, &expert_total); + if (compute_ds4_expert_memory_info(w_, &moe_placement_, info, &memory_error)) { + const uint64_t reserve = std::max( + 2ULL * 1024 * 1024 * 1024, (uint64_t) expert_total / 20); + const uint64_t usable = expert_free > reserve + ? (uint64_t) expert_free - reserve : 0; + stream_cold = info.cold_bytes > usable; + if (stream_cold) nvme_device_cache_bytes = (size_t) usable; + std::fprintf(stderr, + "[deepseek4] Strix cold tier: cold=%.2f GiB free=%.2f GiB " + "reserve=%.2f GiB warm-cache=%.2f GiB mode=%s\n", + gib(info.cold_bytes), gib(expert_free), gib(reserve), + gib(nvme_device_cache_bytes), + stream_cold ? "ssd-stream" : "resident"); + } + } else if (std::strcmp(nvme_mode, "0") != 0 && + std::strcmp(nvme_mode, "off") != 0 && + std::strcmp(nvme_mode, "false") != 0) { + std::fprintf(stderr, + "[deepseek4] ignoring invalid DFLASH_MOE_NVME_COLD_TIER=%s\n", + nvme_mode); + } + hybrid_cfg.materialize_cold_experts = !stream_cold; } if (!build_deepseek4_moe_hybrid_storage_from_file_with_mmap( cfg_.model_path, backend_, w_, moe_placement_, &hybrid_cfg, @@ -901,15 +946,25 @@ bool DeepSeek4Backend::init_hybrid_model() { std::fprintf(stderr, "[deepseek4] failed to compute streaming expert size\n"); return false; } - if (!stream_engine_.init(backend_, max_expert_bytes, &err)) { + ggml_backend_t stream_backend = expert_backend_ ? expert_backend_ : backend_; + MoeStreamConfig stream_config = MoeStreamConfig::from_env(); + if (!std::getenv("DFLASH_MOE_NVME_DEVICE_CACHE_MB") && + nvme_device_cache_bytes > 0) { + stream_config.device_cache_bytes = nvme_device_cache_bytes; + } + if (!stream_engine_.init( + stream_backend, max_expert_bytes, *hybrid, stream_config, &err)) { std::fprintf(stderr, "[deepseek4] failed to init cold-expert stream engine: %s\n", err.c_str()); return false; } std::fprintf(stderr, - "[deepseek4] cold-expert stream engine ready: pinned=%.1f MiB scratch=%.1f MiB\n", + "[deepseek4] cold-expert SSD engine ready: io=%s pinned=%.1f MiB " + "strix_cache=%.1f MiB slots=%d\n", + stream_engine_.io_backend_name(), stream_engine_.pinned_bytes() / 1024.0 / 1024.0, - stream_engine_.scratch_bytes() / 1024.0 / 1024.0); + stream_engine_.device_cache_bytes() / 1024.0 / 1024.0, + stream_engine_.device_slot_count()); } moe_hybrid_ = std::move(hybrid); @@ -1131,7 +1186,8 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, spec_hooks.capture_out = &spec_cap; hp = &spec_hooks; } - if (moe_hybrid_ && (expert_runtime_.compute || expert_backend_)) { + if (moe_hybrid_ && !stream_engine_.is_bound() && + (expert_runtime_.compute || expert_backend_)) { ok = deepseek4_step_layer_range( backend_, cfg_.device.gpu, w_, cache_, hc_state, embed.data(), n_tok, pos, @@ -1234,7 +1290,8 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, const int pos = std::max(0, committed + generated - 1); bool ok = false; - if (moe_hybrid_ && (expert_runtime_.compute || expert_backend_)) { + if (moe_hybrid_ && !stream_engine_.is_bound() && + (expert_runtime_.compute || expert_backend_)) { std::vector hc_state; ok = deepseek4_step_layer_range( backend_, cfg_.device.gpu, w_, cache_, hc_state, diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index ee212bb5b..521248b94 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -2950,7 +2950,7 @@ static bool eval_ds4_hybrid( hybrid_owner->mmap_data, hybrid_owner->mmap_size, hybrid_cfg, desc, regions, storage, token_inp, token_selected, token_weights, 1, - cold_out, &err)) { + cold_out, &err, layer)) { std::fprintf(stderr, "[deepseek4] layer %d cold streaming eval failed: %s\n", layer, err.c_str()); diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index aa4256430..0cac75821 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -992,8 +992,6 @@ bool build_deepseek4_moe_hybrid_storage_from_file_with_mmap( if (err) *err = mmap_err; return false; } - mmap.close_fd(); - const size_t data_start = gguf_get_data_offset(gctx); const auto * file_bytes = static_cast(mmap.addr); std::vector layer_file_data((size_t)w.n_layer); @@ -1039,7 +1037,14 @@ bool build_deepseek4_moe_hybrid_storage_from_file_with_mmap( const MoeHybridConfig cfg = cfg_override ? *cfg_override : make_ds4_moe_hybrid_config(w); const bool ok = build_moe_hybrid_storage_from_file_with_mmap( cfg, backend, placement, layer_descs, layer_file_data, - mmap.addr, mmap.len, out, err, 0, cold_gpu_backend); + mmap.addr, mmap.len, out, err, 0, cold_gpu_backend, +#if defined(_WIN32) + -1 +#else + mmap.fd +#endif + ); + mmap.close_fd(); if (!ok) { mmap.close_map(); diff --git a/server/src/laguna/laguna_backend.cpp b/server/src/laguna/laguna_backend.cpp index db2d366e0..8f600111f 100644 --- a/server/src/laguna/laguna_backend.cpp +++ b/server/src/laguna/laguna_backend.cpp @@ -2071,8 +2071,9 @@ bool LagunaBackend::init_hybrid_mode() { if (eb > max_expert_bytes) max_expert_bytes = eb; } std::string stream_err; - if (stream_engine_.init(backend_, max_expert_bytes, &stream_err)) { - std::printf("[laguna-hybrid] stream engine ready: pinned=%.1f MiB scratch=%.1f MiB\n", + if (stream_engine_.init(backend_, max_expert_bytes, *moe_hybrid_, &stream_err)) { + std::printf("[laguna-hybrid] SSD stream engine ready: io=%s pinned=%.1f MiB scratch=%.1f MiB\n", + stream_engine_.io_backend_name(), stream_engine_.pinned_bytes() / 1024.0 / 1024.0, stream_engine_.scratch_bytes() / 1024.0 / 1024.0); } else { @@ -3129,14 +3130,10 @@ bool LagunaBackend::build_hybrid_storage_from_file( } const size_t file_size = _mf.size(); // Transfer mmap ownership out of the RAII wrapper: the hybrid storage keeps - // the mapping alive and unmaps it in ~MoeHybridStorage. On POSIX the fd can - // be closed now (the mapping stays valid); on Windows release() already - // closed the mapping handle. + // the mapping alive and unmaps it in ~MoeHybridStorage. Keep the POSIX fd + // through storage construction so it can retain an async-I/O duplicate. GgufMmap::OwnedRegion _region = _mf.release(); const void * mmap_addr = _region.data; -#if !defined(_WIN32) - if (_region.fd >= 0) ::close(_region.fd); -#endif const size_t data_start = gguf_get_data_offset(gctx); const auto * file_bytes = (const uint8_t *)mmap_addr; @@ -3168,9 +3165,12 @@ bool LagunaBackend::build_hybrid_storage_from_file( int cache_slots = 0; if (const char * cs = std::getenv("DFLASH_LAGUNA_CACHE_SLOTS")) cache_slots = std::max(0, std::atoi(cs)); else if (cache_slots_ >= 0) cache_slots = cache_slots_; - bool ok = build_moe_hybrid_storage_from_file_with_mmap(hybrid_cfg, backend_, placement, - layer_descs, layer_file_data, - mmap_addr, file_size, *hybrid, &err, cache_slots); + bool ok = build_moe_hybrid_storage_from_file_with_mmap( + hybrid_cfg, backend_, placement, layer_descs, layer_file_data, + mmap_addr, file_size, *hybrid, &err, cache_slots, nullptr, _region.fd); +#if !defined(_WIN32) + if (_region.fd >= 0) ::close(_region.fd); +#endif gguf_free(gctx); if (!ok) { #if defined(_WIN32) diff --git a/server/src/qwen35moe/qwen35moe_backend.cpp b/server/src/qwen35moe/qwen35moe_backend.cpp index 78052a757..d7c750a58 100644 --- a/server/src/qwen35moe/qwen35moe_backend.cpp +++ b/server/src/qwen35moe/qwen35moe_backend.cpp @@ -170,13 +170,10 @@ bool Qwen35MoeBackend::load_target_model(ggml_backend_t backend, TargetWeights & const size_t file_size = _mf.size(); // Transfer mmap ownership out of the RAII wrapper: the hybrid storage // keeps the mapping alive for streaming prefill and unmaps it in - // ~MoeHybridStorage. On POSIX the fd can be closed now (the mapping - // stays valid); on Windows release() already closed the mapping handle. + // ~MoeHybridStorage. Keep the POSIX fd through storage construction so + // it can retain a duplicate for io_uring/pread, then close our copy. GgufMmap::OwnedRegion _region = _mf.release(); const void * mmap_addr = _region.data; -#if !defined(_WIN32) - if (_region.fd >= 0) ::close(_region.fd); -#endif const size_t data_start = gguf_get_data_offset(gctx); const auto * file_bytes = (const uint8_t *)mmap_addr; @@ -214,7 +211,13 @@ bool Qwen35MoeBackend::load_target_model(ggml_backend_t backend, TargetWeights & int cache_slots = 0; if (const char * cs = std::getenv("DFLASH_QWEN35MOE_CACHE_SLOTS")) cache_slots = std::max(0, std::atoi(cs)); else if (cache_slots_ >= 0) cache_slots = cache_slots_; - if (!build_moe_hybrid_storage_from_file_with_mmap(hybrid_cfg, backend, placement, layer_descs, layer_file_data, mmap_addr, file_size, *hybrid, &err, cache_slots)) { + const bool hybrid_ok = build_moe_hybrid_storage_from_file_with_mmap( + hybrid_cfg, backend, placement, layer_descs, layer_file_data, + mmap_addr, file_size, *hybrid, &err, cache_slots, nullptr, _region.fd); +#if !defined(_WIN32) + if (_region.fd >= 0) ::close(_region.fd); +#endif + if (!hybrid_ok) { #if defined(_WIN32) UnmapViewOfFile(const_cast(mmap_addr)); #else @@ -244,8 +247,9 @@ bool Qwen35MoeBackend::load_target_model(ggml_backend_t backend, TargetWeights & } if (max_expert_bytes > 0) { std::string stream_err; - if (stream_engine_.init(backend, max_expert_bytes, &stream_err)) { - std::printf("[qwen35moe] streaming prefill engine ready (pinned=%.1f MiB, scratch=%.1f MiB)\n", + if (stream_engine_.init(backend, max_expert_bytes, *out.moe_hybrid, &stream_err)) { + std::printf("[qwen35moe] SSD stream engine ready (io=%s pinned=%.1f MiB, scratch=%.1f MiB)\n", + stream_engine_.io_backend_name(), stream_engine_.pinned_bytes() / 1024.0 / 1024.0, stream_engine_.scratch_bytes() / 1024.0 / 1024.0); } else { diff --git a/server/test/bench_kimi_k3_hetero.cpp b/server/test/bench_kimi_k3_hetero.cpp new file mode 100644 index 000000000..bfc98bd5d --- /dev/null +++ b/server/test/bench_kimi_k3_hetero.cpp @@ -0,0 +1,386 @@ +// Read-only Kimi-K3 routed-core qualification for heterogeneous Lucebox. +// +// This is not a substitute for the Kimi model graph. It isolates the part +// whose placement is genuinely new: exact routed experts moving from NVMe to +// one GPU while that GPU evaluates the persistent IQ1_S + SiTU expert graph. +// The source file only supplies bytes and is never modified. + +#include "common/moe_hybrid_stream.h" + +#include "ggml-backend.h" +#include "ggml-cuda.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#include +#include +#else +#include +#include +#include +#endif + +using namespace dflash::common; + +namespace { + +constexpr int kKimiExperts = 896; +constexpr int kKimiTopK = 16; +constexpr int kKimiMoeLayers = 92; +constexpr int64_t kKimiLatent = 3584; +constexpr int64_t kKimiExpertFf = 3072; +constexpr float kSituBeta = 4.0f; +constexpr float kSituLinearBeta = 25.0f; + +bool parse_nonnegative(const char * text, uint64_t & out) { + if (!text || !text[0] || text[0] == '-') return false; + char * end = nullptr; + errno = 0; + const unsigned long long value = std::strtoull(text, &end, 10); + if (errno != 0 || end == text || *end != '\0') return false; + out = (uint64_t) value; + return true; +} + +double gib(uint64_t bytes) { + return (double) bytes / (1024.0 * 1024.0 * 1024.0); +} + +class PersistentKimiExpertGraph { +public: + ~PersistentKimiExpertGraph() { destroy(); } + + bool init(ggml_backend_t backend, std::string & error) { + backend_ = backend; + ggml_init_params params{}; + params.mem_size = 16 * 1024 * 1024; + params.no_alloc = true; + ctx_ = ggml_init(params); + if (!ctx_) { + error = "ggml_init failed for Kimi expert graph"; + return false; + } + + gate_ = ggml_new_tensor_2d( + ctx_, GGML_TYPE_IQ1_S, kKimiLatent, kKimiExpertFf); + up_ = ggml_new_tensor_2d( + ctx_, GGML_TYPE_IQ1_S, kKimiLatent, kKimiExpertFf); + down_ = ggml_new_tensor_2d( + ctx_, GGML_TYPE_IQ1_S, kKimiExpertFf, kKimiLatent); + input_ = ggml_new_tensor_2d(ctx_, GGML_TYPE_F32, kKimiLatent, 1); + ggml_set_input(gate_); + ggml_set_input(up_); + ggml_set_input(down_); + ggml_set_input(input_); + + ggml_tensor * gate_value = ggml_mul_mat(ctx_, gate_, input_); + ggml_tensor * up_value = ggml_mul_mat(ctx_, up_, input_); + + // SiTU(g, u) = beta*tanh(g/beta)*sigmoid(g) + // * linear_beta*tanh(u/linear_beta). + ggml_tensor * activated = ggml_scale(ctx_, gate_value, 1.0f / kSituBeta); + activated = ggml_tanh(ctx_, activated); + activated = ggml_scale(ctx_, activated, kSituBeta); + activated = ggml_mul(ctx_, activated, ggml_sigmoid(ctx_, gate_value)); + up_value = ggml_scale(ctx_, up_value, 1.0f / kSituLinearBeta); + up_value = ggml_tanh(ctx_, up_value); + up_value = ggml_scale(ctx_, up_value, kSituLinearBeta); + activated = ggml_mul(ctx_, activated, up_value); + output_ = ggml_mul_mat(ctx_, down_, activated); + ggml_set_output(output_); + + graph_ = ggml_new_graph_custom(ctx_, 256, false); + ggml_build_forward_expand(graph_, output_); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_, backend_); + if (!buffer_) { + error = "device allocation failed for Kimi expert graph"; + return false; + } + + std::vector input((size_t) kKimiLatent); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = 0.01f * std::sin((float) i * 0.013f); + } + ggml_backend_tensor_set( + input_, input.data(), 0, input.size() * sizeof(float)); + return true; + } + + bool launch(const MoeHybridStreamEngine & engine, std::string & error) { + if (!ctx_ || !graph_ || !backend_) { + error = "Kimi expert graph is not initialized"; + return false; + } + if (engine.scratch_gate_bytes() != ggml_nbytes(gate_) || + engine.scratch_up_bytes() != ggml_nbytes(up_) || + engine.scratch_down_bytes() != ggml_nbytes(down_)) { + error = "streamed Kimi expert byte layout does not match IQ1_S graph"; + return false; + } + gate_->data = const_cast(engine.scratch_gate_data()); + up_->data = const_cast(engine.scratch_up_data()); + down_->data = const_cast(engine.scratch_down_data()); + if (ggml_backend_graph_compute_async(backend_, graph_) != + GGML_STATUS_SUCCESS) { + error = "Kimi expert graph launch failed"; + return false; + } + return true; + } + + void destroy() { + if (backend_) ggml_backend_synchronize(backend_); + if (buffer_) { + ggml_backend_buffer_free(buffer_); + buffer_ = nullptr; + } + if (ctx_) { + ggml_free(ctx_); + ctx_ = nullptr; + } + graph_ = nullptr; + backend_ = nullptr; + gate_ = up_ = down_ = input_ = output_ = nullptr; + } + +private: + ggml_backend_t backend_ = nullptr; + ggml_context * ctx_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; + ggml_tensor * gate_ = nullptr; + ggml_tensor * up_ = nullptr; + ggml_tensor * down_ = nullptr; + ggml_tensor * input_ = nullptr; + ggml_tensor * output_ = nullptr; +}; + +std::vector route_for(int token, int layer, int top_k, bool repeat) { + std::vector population((size_t) kKimiExperts); + std::iota(population.begin(), population.end(), 0); + const uint64_t token_key = repeat ? 0 : (uint64_t) token; + std::mt19937_64 rng( + 0x4b494d49334c5543ULL ^ (token_key * 0x9e3779b97f4a7c15ULL) ^ + ((uint64_t) layer * 0xbf58476d1ce4e5b9ULL)); + for (int i = 0; i < top_k; ++i) { + std::uniform_int_distribution choose(i, kKimiExperts - 1); + const int selected = choose(rng); + std::swap(population[(size_t) i], population[(size_t) selected]); + } + population.resize((size_t) top_k); + return population; +} + +} // namespace + +int main(int argc, char ** argv) { + if (argc < 2 || argc > 8) { + std::fprintf(stderr, + "usage: %s MODEL_FILE [device=1] [tokens=1] [layers=92] " + "[top_k=16] [compute=1] [repeat_routes=0]\n", + argv[0]); + return 2; + } + + uint64_t device_arg = 1; + uint64_t tokens_arg = 1; + uint64_t layers_arg = kKimiMoeLayers; + uint64_t top_k_arg = kKimiTopK; + uint64_t compute_arg = 1; + uint64_t repeat_arg = 0; + uint64_t * values[] = { + &device_arg, &tokens_arg, &layers_arg, &top_k_arg, + &compute_arg, &repeat_arg, + }; + for (int i = 2; i < argc; ++i) { + if (!parse_nonnegative(argv[i], *values[i - 2])) { + std::fprintf(stderr, "invalid integer: %s\n", argv[i]); + return 2; + } + } + if (tokens_arg == 0 || layers_arg == 0 || layers_arg > kKimiMoeLayers || + top_k_arg == 0 || top_k_arg > kKimiExperts || compute_arg > 1 || + repeat_arg > 1) { + std::fprintf(stderr, "Kimi benchmark arguments are out of range\n"); + return 2; + } + if (device_arg >= (uint64_t) ggml_backend_cuda_get_device_count()) { + std::fprintf(stderr, "GPU device is out of range\n"); + return 2; + } + +#if defined(_WIN32) + const int fd = ::_open(argv[1], _O_RDONLY | _O_BINARY); + struct _stat64 st{}; + if (fd < 0 || ::_fstat64(fd, &st) != 0 || st.st_size <= 0) { +#else + const int fd = ::open(argv[1], O_RDONLY | O_CLOEXEC); + struct stat st{}; + if (fd < 0 || ::fstat(fd, &st) != 0 || st.st_size <= 0) { +#endif + std::fprintf(stderr, "cannot open input file: %s\n", std::strerror(errno)); + return 1; + } + const size_t file_bytes = (size_t) st.st_size; + + const size_t gate_bytes = + ggml_row_size(GGML_TYPE_IQ1_S, kKimiLatent) * (size_t) kKimiExpertFf; + const size_t up_bytes = gate_bytes; + const size_t down_bytes = + ggml_row_size(GGML_TYPE_IQ1_S, kKimiExpertFf) * (size_t) kKimiLatent; + const size_t expert_bytes = gate_bytes + up_bytes + down_bytes; + const size_t gate_stack = gate_bytes * (size_t) kKimiExperts; + const size_t up_stack = up_bytes * (size_t) kKimiExperts; + const size_t down_stack = down_bytes * (size_t) kKimiExperts; + const size_t required_bytes = gate_stack + up_stack + down_stack; + if (file_bytes < required_bytes) { + std::fprintf(stderr, + "input needs at least %.3f GiB for one Kimi expert stack\n", + gib(required_bytes)); + return 2; + } + + LayerExpertRegions one_layer; + one_layer.fused_gate_up = false; + one_layer.expert_bytes_gate = gate_bytes; + one_layer.expert_bytes_up = up_bytes; + one_layer.expert_bytes_down = down_bytes; + one_layer.gate_exps = {0, gate_stack}; + one_layer.up_exps = {gate_stack, up_stack}; + one_layer.down_exps = {gate_stack + up_stack, down_stack}; + std::vector regions((size_t) layers_arg, one_layer); + + ggml_backend_t backend = ggml_backend_cuda_init((int) device_arg); + if (!backend) { + std::fprintf(stderr, "failed to initialize GPU backend\n"); + return 1; + } + + MoeHybridStorage storage; + storage.mmap_size = file_bytes; +#if defined(_WIN32) + storage.mmap_fd = -1; +#else + storage.mmap_fd = ::dup(fd); +#endif + storage.layer_regions = regions; + + MoeStreamConfig stream_config = MoeStreamConfig::from_env(); + MoeHybridStreamEngine engine; + std::string error; + if (!engine.init(backend, expert_bytes, storage, stream_config, &error)) { + std::fprintf(stderr, "stream engine initialization failed: %s\n", error.c_str()); + ggml_backend_free(backend); + return 1; + } + + PersistentKimiExpertGraph expert_graph; + if (compute_arg && !expert_graph.init(backend, error)) { + std::fprintf(stderr, "Kimi graph initialization failed: %s\n", error.c_str()); + engine.destroy(); + ggml_backend_free(backend); + return 1; + } + + const uint64_t accesses = + tokens_arg * layers_arg * top_k_arg; + const auto begin = std::chrono::steady_clock::now(); + for (int token = 0; token < (int) tokens_arg; ++token) { + for (int layer = 0; layer < (int) layers_arg; ++layer) { + const std::vector experts = route_for( + token, layer, (int) top_k_arg, repeat_arg != 0); + engine.request_experts( + layer, experts.data(), (int) experts.size(), MoeNvmePriority::Demand); + + int staged_slot = -1; + if (!engine.stage_expert_cached_async( + layer, experts[0], &staged_slot, &error)) { + std::fprintf(stderr, "initial expert stage failed: %s\n", error.c_str()); + return 1; + } + for (size_t i = 0; i < experts.size(); ++i) { + const int current_slot = staged_slot; + if (!engine.activate_device_slot(current_slot, &error)) { + std::fprintf(stderr, "expert activation failed: %s\n", error.c_str()); + return 1; + } + if (compute_arg && !expert_graph.launch(engine, error)) { + std::fprintf(stderr, "expert compute failed: %s\n", error.c_str()); + return 1; + } + + // Disk/H2D for expert N+1 overlaps expert N's persistent graph. + if (i + 1 < experts.size()) { + int next_slot = -1; + if (!engine.stage_expert_cached_async( + layer, experts[i + 1], &next_slot, &error)) { + ggml_backend_synchronize(backend); + std::fprintf(stderr, "pipelined expert stage failed: %s\n", error.c_str()); + return 1; + } + staged_slot = next_slot; + } + if (compute_arg) ggml_backend_synchronize(backend); + engine.release_device_slot(current_slot); + } + } + } + ggml_backend_synchronize(backend); + const auto end = std::chrono::steady_clock::now(); + const double seconds = + std::chrono::duration(end - begin).count(); + const MoeNvmeStats stats = engine.io_stats(); + const double misses = expert_bytes > 0 + ? (double) stats.payload_bytes / (double) expert_bytes : 0.0; + const double hit_rate = accesses > 0 + ? std::max(0.0, 1.0 - misses / (double) accesses) : 0.0; + char description[256] = {}; + ggml_backend_cuda_get_device_description( + (int) device_arg, description, sizeof(description)); + + std::printf( + "kimi_k3 device=%" PRIu64 " description=%s latent=%lld ff=%lld " + "experts=%d top_k=%" PRIu64 " layers=%" PRIu64 " tokens=%" PRIu64 + " compute=%s repeat_routes=%s\n", + device_arg, description, (long long) kKimiLatent, + (long long) kKimiExpertFf, kKimiExperts, top_k_arg, layers_arg, + tokens_arg, compute_arg ? "situ-iq1s" : "off", + repeat_arg ? "yes" : "no"); + std::printf( + "expert_bytes=%zu expert_mib=%.6f accesses=%" PRIu64 + " elapsed_s=%.6f routed_core_tok_s=%.6f experts_s=%.2f\n", + expert_bytes, expert_bytes / 1024.0 / 1024.0, accesses, seconds, + seconds > 0 ? (double) tokens_arg / seconds : 0.0, + seconds > 0 ? (double) accesses / seconds : 0.0); + std::printf( + "ssd_payload_gib=%.6f physical_gib=%.6f pipeline_gib_s=%.6f " + "estimated_device_cache_hit=%.4f cache_gib=%.3f io_errors=%" PRIu64 "\n", + gib(stats.payload_bytes), gib(stats.physical_bytes), + seconds > 0 ? gib(stats.payload_bytes) / seconds : 0.0, + hit_rate, gib(engine.device_cache_bytes()), stats.errors); + + expert_graph.destroy(); + engine.destroy(); + ggml_backend_free(backend); +#if defined(_WIN32) + ::_close(fd); +#else + ::close(fd); +#endif + return stats.errors == 0 ? 0 : 1; +} diff --git a/server/test/bench_moe_nvme_io.cpp b/server/test/bench_moe_nvme_io.cpp new file mode 100644 index 000000000..feb0234ea --- /dev/null +++ b/server/test/bench_moe_nvme_io.cpp @@ -0,0 +1,241 @@ +// Read-only SSD benchmark for the routed-MoE scheduler. +// +// Usage: +// bench_moe_nvme_io MODEL_FILE [expert_mib] [working_set] [rounds] [batch] +// +// DFLASH_MOE_NVME_BACKEND and DFLASH_MOE_NVME_DIRECT select the same path as +// production. The benchmark treats two large, disjoint ranges of the input as +// fused gate/up and down expert tensors. It never modifies the input file. + +#include "common/moe_nvme_scheduler.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#include +#include +#include +#else +#include +#include +#include +#endif + +using namespace dflash::common; + +namespace { + +bool aligned_allocate(void ** ptr, size_t bytes, void *) { +#if defined(_WIN32) + *ptr = _aligned_malloc(bytes, 4096); + return *ptr != nullptr; +#else + return ::posix_memalign(ptr, 4096, bytes) == 0; +#endif +} + +void aligned_free(void * ptr, void *) { +#if defined(_WIN32) + _aligned_free(ptr); +#else + std::free(ptr); +#endif +} + +bool parse_positive(const char * text, uint64_t & out) { + if (!text || !text[0] || text[0] == '-') return false; + char * end = nullptr; + errno = 0; + const unsigned long long value = std::strtoull(text, &end, 10); + if (errno != 0 || end == text || *end != '\0' || value == 0) return false; + out = (uint64_t) value; + return true; +} + +size_t align_down(size_t value, size_t alignment) { + return value & ~(alignment - 1); +} + +double gib(uint64_t bytes) { + return (double) bytes / (1024.0 * 1024.0 * 1024.0); +} + +} // namespace + +int main(int argc, char ** argv) { + if (argc < 2 || argc > 6) { + std::fprintf(stderr, + "usage: %s MODEL_FILE [expert_mib=24] [working_set=64] " + "[rounds=4] [batch=8]\n", argv[0]); + return 2; + } + + uint64_t expert_mib = 24; + uint64_t requested_working_set = 64; + uint64_t rounds = 4; + uint64_t requested_batch = 8; + uint64_t * values[] = {&expert_mib, &requested_working_set, &rounds, + &requested_batch}; + for (int i = 2; i < argc; ++i) { + if (!parse_positive(argv[i], *values[i - 2])) { + std::fprintf(stderr, "invalid positive integer: %s\n", argv[i]); + return 2; + } + } + if (expert_mib > (uint64_t) std::numeric_limits::max() / (1024 * 1024)) { + std::fprintf(stderr, "expert size is too large\n"); + return 2; + } + +#if defined(_WIN32) + const int fd = ::_open(argv[1], _O_RDONLY | _O_BINARY); + struct _stat64 st{}; + if (fd < 0 || ::_fstat64(fd, &st) != 0 || st.st_size <= 0) { +#else + const int fd = ::open(argv[1], O_RDONLY | O_CLOEXEC); + struct stat st{}; + if (fd < 0 || ::fstat(fd, &st) != 0 || st.st_size <= 0) { +#endif + std::fprintf(stderr, "cannot open input file: %s\n", std::strerror(errno)); + if (fd >= 0) { +#if defined(_WIN32) + ::_close(fd); +#else + ::close(fd); +#endif + } + return 1; + } + + const size_t file_bytes = (size_t) st.st_size; + const size_t alignment = 4096; + const size_t requested_payload = (size_t) expert_mib * 1024 * 1024; + // Keep both tensor slices naturally aligned so direct-I/O throughput is + // measured rather than an artificial alignment-amplification worst case. + const size_t gate_up_bytes = align_down(requested_payload * 2 / 3, alignment); + const size_t down_bytes = align_down(requested_payload - gate_up_bytes, alignment); + const size_t payload_bytes = gate_up_bytes + down_bytes; + if (gate_up_bytes == 0 || down_bytes == 0 || payload_bytes == 0) { + std::fprintf(stderr, "expert payload is too small\n"); + return 2; + } + + const size_t possible_experts = file_bytes / payload_bytes; + const size_t working_set = std::min( + (size_t) requested_working_set, possible_experts); + if (working_set < 2) { + std::fprintf(stderr, "file is too small for two %.1f MiB experts\n", + payload_bytes / 1024.0 / 1024.0); + return 2; + } + + LayerExpertRegions region; + region.fused_gate_up = true; + region.expert_bytes_gate_up = gate_up_bytes; + region.expert_bytes_down = down_bytes; + region.gate_up_exps = {0, gate_up_bytes * working_set}; + region.down_exps = {region.gate_up_exps.size, down_bytes * working_set}; + + MoeNvmeConfig config = MoeNvmeConfig::from_env(); + const size_t batch = std::max(1, std::min( + {(size_t) requested_batch, working_set, (size_t) config.host_slots})); + MoeNvmeScheduler scheduler; + std::string error; + if (!scheduler.init(config, payload_bytes, aligned_allocate, aligned_free, + nullptr, &error) || + !scheduler.bind_source({nullptr, file_bytes, fd}, {region}, &error)) { + std::fprintf(stderr, "scheduler initialization failed: %s\n", error.c_str()); + return 1; + } + +#if defined(__linux__) + (void) ::posix_fadvise(fd, 0, 0, POSIX_FADV_RANDOM); +#endif + + const uint64_t total_requests = rounds * working_set; + uint64_t completed = 0; + uint64_t checksum = 0; + std::vector batch_ms; + batch_ms.reserve((size_t) ((total_requests + batch - 1) / batch)); + const auto wall_begin = std::chrono::steady_clock::now(); + while (completed < total_requests) { + const size_t count = (size_t) std::min(batch, total_requests - completed); + std::vector experts(count); + for (size_t i = 0; i < count; ++i) { + // Rotate every round and walk the full working set. This avoids + // measuring a tiny cache-resident subset while remaining repeatable. + const uint64_t ordinal = completed + i; + const uint64_t round = ordinal / working_set; + const uint64_t within = ordinal % working_set; + experts[i] = (int) ((within + round * (working_set / 2 + 1)) % working_set); + } + + const auto batch_begin = std::chrono::steady_clock::now(); + for (int expert : experts) { + if (!scheduler.request(0, expert, MoeNvmePriority::Demand, &error)) { + std::fprintf(stderr, "request failed: %s\n", error.c_str()); + return 1; + } + } + for (int expert : experts) { + MoeNvmeLease lease; + if (!scheduler.acquire(0, expert, lease, &error)) { + std::fprintf(stderr, "acquire failed: %s\n", error.c_str()); + return 1; + } + for (int span = 0; span < lease.layout().span_count; ++span) { + const MoeExpertIoSpan & io = lease.layout().spans[span]; + const uint8_t * data = lease.data() + io.buffer_offset; + checksum += data[0]; + checksum += data[io.bytes - 1]; + } + } + const auto batch_end = std::chrono::steady_clock::now(); + batch_ms.push_back(std::chrono::duration( + batch_end - batch_begin).count()); + completed += count; + } + const auto wall_end = std::chrono::steady_clock::now(); + const double seconds = std::chrono::duration(wall_end - wall_begin).count(); + const MoeNvmeStats stats = scheduler.stats(); + std::sort(batch_ms.begin(), batch_ms.end()); + const size_t p50_index = (batch_ms.size() - 1) / 2; + const size_t p95_index = (batch_ms.size() - 1) * 95 / 100; + + std::printf("backend=%s direct=%s file_gib=%.2f expert_mib=%.2f " + "working_set=%zu rounds=%" PRIu64 " batch=%zu slots=%d\n", + scheduler.effective_backend_name(), + scheduler.direct_io_active() ? "yes" : "no", gib(file_bytes), + payload_bytes / 1024.0 / 1024.0, working_set, rounds, batch, + scheduler.slot_count()); + std::printf("elapsed_s=%.6f payload_gib=%.3f physical_gib=%.3f " + "payload_gib_s=%.3f physical_gib_s=%.3f experts_s=%.2f\n", + seconds, gib(stats.payload_bytes), gib(stats.physical_bytes), + seconds > 0 ? gib(stats.payload_bytes) / seconds : 0.0, + seconds > 0 ? gib(stats.physical_bytes) / seconds : 0.0, + seconds > 0 ? stats.payload_bytes / (double) payload_bytes / seconds : 0.0); + std::printf("batch_ms_p50=%.3f batch_ms_p95=%.3f cache_hits=%" PRIu64 + " dedupe=%" PRIu64 " evictions=%" PRIu64 + " errors=%" PRIu64 " checksum=%" PRIu64 "\n", + batch_ms[p50_index], batch_ms[p95_index], stats.cache_hits, + stats.inflight_deduplications, stats.evictions, stats.errors, + checksum); + + scheduler.destroy(); +#if defined(_WIN32) + ::_close(fd); +#else + ::close(fd); +#endif + return stats.errors == 0 ? 0 : 1; +} diff --git a/server/test/bench_moe_nvme_pipeline.cpp b/server/test/bench_moe_nvme_pipeline.cpp new file mode 100644 index 000000000..3c22452ca --- /dev/null +++ b/server/test/bench_moe_nvme_pipeline.cpp @@ -0,0 +1,223 @@ +// Read-only end-to-end NVMe -> pinned host -> HIP/CUDA device benchmark. +// It exercises MoeHybridStreamEngine itself on either Lucebox GPU. + +#include "common/moe_hybrid_stream.h" + +#include "ggml-backend.h" +#include "ggml-cuda.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#include +#include +#else +#include +#include +#include +#endif + +using namespace dflash::common; + +namespace { + +bool parse_positive(const char * text, uint64_t & out) { + if (!text || !text[0] || text[0] == '-') return false; + char * end = nullptr; + errno = 0; + const unsigned long long value = std::strtoull(text, &end, 10); + if (errno != 0 || end == text || *end != '\0' || value == 0) return false; + out = (uint64_t) value; + return true; +} + +size_t align_down(size_t value, size_t alignment) { + return value & ~(alignment - 1); +} + +double gib(uint64_t bytes) { + return (double) bytes / (1024.0 * 1024.0 * 1024.0); +} + +} // namespace + +int main(int argc, char ** argv) { + if (argc < 2 || argc > 7) { + std::fprintf(stderr, + "usage: %s MODEL_FILE [device=0] [expert_mib=24] " + "[working_set=64] [rounds=4] [batch=8]\n", argv[0]); + return 2; + } + + uint64_t device_arg = 0; + uint64_t expert_mib = 24; + uint64_t requested_working_set = 64; + uint64_t rounds = 4; + uint64_t requested_batch = 8; + uint64_t * values[] = {&device_arg, &expert_mib, &requested_working_set, + &rounds, &requested_batch}; + for (int i = 2; i < argc; ++i) { + // Device zero is valid; other arguments must be positive. + if (i == 2 && std::strcmp(argv[i], "0") == 0) continue; + if (!parse_positive(argv[i], *values[i - 2])) { + std::fprintf(stderr, "invalid integer: %s\n", argv[i]); + return 2; + } + } + if (device_arg >= (uint64_t) ggml_backend_cuda_get_device_count()) { + std::fprintf(stderr, "GPU device is out of range\n"); + return 2; + } + if (expert_mib > (uint64_t) std::numeric_limits::max() / (1024 * 1024)) { + std::fprintf(stderr, "expert size is too large\n"); + return 2; + } + +#if defined(_WIN32) + const int fd = ::_open(argv[1], _O_RDONLY | _O_BINARY); + struct _stat64 st{}; + if (fd < 0 || ::_fstat64(fd, &st) != 0 || st.st_size <= 0) { +#else + const int fd = ::open(argv[1], O_RDONLY | O_CLOEXEC); + struct stat st{}; + if (fd < 0 || ::fstat(fd, &st) != 0 || st.st_size <= 0) { +#endif + std::fprintf(stderr, "cannot open input file: %s\n", std::strerror(errno)); + return 1; + } + + const size_t file_bytes = (size_t) st.st_size; + const size_t alignment = 4096; + const size_t requested_payload = (size_t) expert_mib * 1024 * 1024; + const size_t gate_up_bytes = align_down(requested_payload * 2 / 3, alignment); + const size_t down_bytes = align_down(requested_payload - gate_up_bytes, alignment); + const size_t payload_bytes = gate_up_bytes + down_bytes; + const size_t possible_experts = payload_bytes ? file_bytes / payload_bytes : 0; + const size_t working_set = std::min( + (size_t) requested_working_set, possible_experts); + if (gate_up_bytes == 0 || down_bytes == 0 || working_set < 2) { + std::fprintf(stderr, "file or expert payload is too small\n"); + return 2; + } + + LayerExpertRegions region; + region.fused_gate_up = true; + region.expert_bytes_gate_up = gate_up_bytes; + region.expert_bytes_down = down_bytes; + region.gate_up_exps = {0, gate_up_bytes * working_set}; + region.down_exps = {region.gate_up_exps.size, down_bytes * working_set}; + + ggml_backend_t backend = ggml_backend_cuda_init((int) device_arg); + if (!backend) { + std::fprintf(stderr, "failed to initialize GPU backend\n"); + return 1; + } + + MoeHybridStorage storage; + storage.mmap_size = file_bytes; +#if defined(_WIN32) + storage.mmap_fd = -1; +#else + storage.mmap_fd = ::dup(fd); +#endif + storage.layer_regions = {region}; + MoeHybridStreamEngine engine; + std::string error; + if (!engine.init(backend, payload_bytes, storage, &error)) { + std::fprintf(stderr, "stream engine initialization failed: %s\n", error.c_str()); + ggml_backend_free(backend); + return 1; + } + + const MoeNvmeConfig io_config = MoeNvmeConfig::from_env(); + const size_t batch = std::max(1, std::min( + {(size_t) requested_batch, working_set, (size_t) io_config.host_slots})); + const uint64_t total_requests = rounds * working_set; + uint64_t completed = 0; + std::vector batch_ms; + batch_ms.reserve((size_t) ((total_requests + batch - 1) / batch)); + const auto wall_begin = std::chrono::steady_clock::now(); + while (completed < total_requests) { + const size_t count = (size_t) std::min(batch, total_requests - completed); + std::vector experts(count); + for (size_t i = 0; i < count; ++i) { + const uint64_t ordinal = completed + i; + const uint64_t round = ordinal / working_set; + const uint64_t within = ordinal % working_set; + experts[i] = (int32_t) ((within + round * (working_set / 2 + 1)) % working_set); + } + const auto batch_begin = std::chrono::steady_clock::now(); + engine.request_experts(0, experts.data(), (int) experts.size(), + MoeNvmePriority::Demand); + int staged_slot = -1; + if (!engine.stage_expert_cached_async(0, experts[0], &staged_slot, &error)) { + std::fprintf(stderr, "initial stage failed: %s\n", error.c_str()); + return 1; + } + for (size_t i = 0; i < count; ++i) { + const int current_slot = staged_slot; + if (i + 1 < count) { + int next_slot = -1; + if (!engine.stage_expert_cached_async( + 0, experts[i + 1], &next_slot, &error)) { + std::fprintf(stderr, "pipelined stage failed: %s\n", error.c_str()); + return 1; + } + staged_slot = next_slot; + } + if (!engine.activate_device_slot(current_slot, &error)) { + std::fprintf(stderr, "activation failed: %s\n", error.c_str()); + return 1; + } + engine.release_device_slot(current_slot); + } + const auto batch_end = std::chrono::steady_clock::now(); + batch_ms.push_back(std::chrono::duration( + batch_end - batch_begin).count()); + completed += count; + } + const auto wall_end = std::chrono::steady_clock::now(); + const double seconds = std::chrono::duration(wall_end - wall_begin).count(); + const MoeNvmeStats stats = engine.io_stats(); + std::sort(batch_ms.begin(), batch_ms.end()); + const size_t p50_index = (batch_ms.size() - 1) / 2; + const size_t p95_index = (batch_ms.size() - 1) * 95 / 100; + char description[256] = {}; + ggml_backend_cuda_get_device_description((int) device_arg, description, + sizeof(description)); + + std::printf("device=%" PRIu64 " description=%s backend=%s " + "expert_mib=%.2f working_set=%zu rounds=%" PRIu64 + " batch=%zu host_mib=%.1f device_mib=%.1f\n", + device_arg, description, engine.io_backend_name(), + payload_bytes / 1024.0 / 1024.0, working_set, rounds, batch, + engine.pinned_bytes() / 1024.0 / 1024.0, + engine.scratch_bytes() / 1024.0 / 1024.0); + std::printf("elapsed_s=%.6f transferred_gib=%.3f pipeline_gib_s=%.3f " + "experts_s=%.2f batch_ms_p50=%.3f batch_ms_p95=%.3f " + "physical_gib=%.3f errors=%" PRIu64 "\n", + seconds, gib(stats.payload_bytes), + seconds > 0 ? gib(stats.payload_bytes) / seconds : 0.0, + seconds > 0 ? total_requests / seconds : 0.0, + batch_ms[p50_index], batch_ms[p95_index], + gib(stats.physical_bytes), stats.errors); + + engine.destroy(); + ggml_backend_free(backend); +#if defined(_WIN32) + ::_close(fd); +#else + ::close(fd); +#endif + return stats.errors == 0 ? 0 : 1; +} diff --git a/server/test/test_moe_nvme_scheduler.cpp b/server/test/test_moe_nvme_scheduler.cpp new file mode 100644 index 000000000..e85014a81 --- /dev/null +++ b/server/test/test_moe_nvme_scheduler.cpp @@ -0,0 +1,364 @@ +#include "CppUnitTestFramework.hpp" +#include "common/moe_nvme_scheduler.h" + +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#else +#include +#include +#endif + +using namespace dflash::common; + +#define NVME_REQUIRE(cond) do { \ + if (!(cond)) throw std::runtime_error(std::string(__FILE__) + ":" + \ + std::to_string(__LINE__) + ": " + #cond); \ +} while (0) + +namespace { + +struct MoeNvmeSchedulerFixture {}; + +bool aligned_allocate(void ** ptr, size_t bytes, void *) { +#if defined(_WIN32) + *ptr = _aligned_malloc(bytes, 4096); + return *ptr != nullptr; +#else + return ::posix_memalign(ptr, 4096, bytes) == 0; +#endif +} + +void aligned_free(void * ptr, void *) { +#if defined(_WIN32) + _aligned_free(ptr); +#else + std::free(ptr); +#endif +} + +uint8_t expected_byte(int layer, int tensor, int expert, size_t offset) { + return (uint8_t) ((layer * 61 + tensor * 31 + expert * 17 + offset * 7) & 0xff); +} + +void fill_tensor(std::vector & file, const ExpertFileRegion & region, + size_t expert_bytes, int layer, int tensor, int experts) { + for (int expert = 0; expert < experts; ++expert) { + for (size_t i = 0; i < expert_bytes; ++i) { + file[region.offset + (size_t) expert * expert_bytes + i] = + expected_byte(layer, tensor, expert, i); + } + } +} + +struct SyntheticModel { + static constexpr int kExperts = 8; + std::vector file; + std::vector regions; + + SyntheticModel() : file(2 * 1024 * 1024, 0xa5), regions(2) { + auto & a = regions[0]; + a.fused_gate_up = false; + a.expert_bytes_gate = 4093; + a.expert_bytes_up = 6141; + a.expert_bytes_down = 8189; + a.gate_exps = {257, a.expert_bytes_gate * kExperts}; + a.up_exps = {a.gate_exps.offset + a.gate_exps.size + 113, + a.expert_bytes_up * kExperts}; + a.down_exps = {a.up_exps.offset + a.up_exps.size + 197, + a.expert_bytes_down * kExperts}; + fill_tensor(file, a.gate_exps, a.expert_bytes_gate, 0, 0, kExperts); + fill_tensor(file, a.up_exps, a.expert_bytes_up, 0, 1, kExperts); + fill_tensor(file, a.down_exps, a.expert_bytes_down, 0, 2, kExperts); + + auto & b = regions[1]; + b.fused_gate_up = true; + b.expert_bytes_gate_up = 10007; + b.expert_bytes_down = 5003; + b.gate_up_exps = {512 * 1024 + 73, b.expert_bytes_gate_up * kExperts}; + b.down_exps = {b.gate_up_exps.offset + b.gate_up_exps.size + 89, + b.expert_bytes_down * kExperts}; + fill_tensor(file, b.gate_up_exps, b.expert_bytes_gate_up, 1, 0, kExperts); + fill_tensor(file, b.down_exps, b.expert_bytes_down, 1, 1, kExperts); + } +}; + +void verify_lease(const MoeNvmeLease & lease, int layer, int expert) { + NVME_REQUIRE(lease); + const MoeExpertIoLayout & layout = lease.layout(); + NVME_REQUIRE(layout.key.layer == layer); + NVME_REQUIRE(layout.key.expert == expert); + const int expected_spans = layer == 0 ? 3 : 2; + NVME_REQUIRE(layout.span_count == expected_spans); + for (int tensor = 0; tensor < layout.span_count; ++tensor) { + const MoeExpertIoSpan & span = layout.spans[tensor]; + const uint8_t * payload = lease.data() + span.buffer_offset; + NVME_REQUIRE(payload[0] == expected_byte(layer, tensor, expert, 0)); + NVME_REQUIRE(payload[span.bytes / 2] == + expected_byte(layer, tensor, expert, span.bytes / 2)); + NVME_REQUIRE(payload[span.bytes - 1] == + expected_byte(layer, tensor, expert, span.bytes - 1)); + if (tensor > 0) { + NVME_REQUIRE(span.device_offset == + layout.spans[tensor - 1].device_offset + + layout.spans[tensor - 1].bytes); + } + NVME_REQUIRE((span.io_file_offset & 4095) == 0); + NVME_REQUIRE((span.io_buffer_offset & 4095) == 0); + NVME_REQUIRE((span.io_bytes & 4095) == 0); + } +} + +} // namespace + +TEST_CASE(MoeNvmeSchedulerFixture, exact_layout_bounds_and_alignment) { + SyntheticModel model; + MoeExpertIoLayout layout; + std::string err; + NVME_REQUIRE(make_moe_expert_io_layout( + 0, 3, model.regions[0], model.file.size(), 4096, layout, &err)); + NVME_REQUIRE(layout.span_count == 3); + NVME_REQUIRE(layout.payload_bytes == + model.regions[0].expert_bytes_gate + + model.regions[0].expert_bytes_up + + model.regions[0].expert_bytes_down); + NVME_REQUIRE(layout.host_bytes >= layout.payload_bytes); + NVME_REQUIRE(!make_moe_expert_io_layout( + 0, SyntheticModel::kExperts, model.regions[0], model.file.size(), + 4096, layout, &err)); + NVME_REQUIRE(!err.empty()); +} + +TEST_CASE(MoeNvmeSchedulerFixture, async_exact_reads_dedupe_and_cache) { + SyntheticModel model; + MoeNvmeConfig config; + config.backend = MoeNvmeBackend::Mmap; + config.direct_io = MoeNvmeDirectMode::Disabled; + config.host_slots = 4; + config.io_threads = 2; + config.demand_reserve = 2; + + MoeNvmeScheduler scheduler; + std::string err; + const size_t max_payload = 32 * 1024; + NVME_REQUIRE(scheduler.init(config, max_payload, aligned_allocate, + aligned_free, nullptr, &err)); + NVME_REQUIRE(scheduler.bind_source( + {model.file.data(), model.file.size(), -1}, model.regions, &err)); + NVME_REQUIRE(std::string(scheduler.effective_backend_name()) == "mmap-workers"); + + NVME_REQUIRE(scheduler.request(0, 3, MoeNvmePriority::Prefetch, &err)); + NVME_REQUIRE(scheduler.request(0, 3, MoeNvmePriority::Prefetch, &err)); + + MoeNvmeLease first; + NVME_REQUIRE(scheduler.acquire(0, 3, first, &err)); + verify_lease(first, 0, 3); + first.reset(); + + MoeNvmeLease cached; + NVME_REQUIRE(scheduler.acquire(0, 3, cached, &err)); + verify_lease(cached, 0, 3); + cached.reset(); + + MoeNvmeLease fused; + NVME_REQUIRE(scheduler.acquire(1, 5, fused, &err)); + verify_lease(fused, 1, 5); + fused.reset(); + + const MoeNvmeStats stats = scheduler.stats(); + NVME_REQUIRE(stats.requests >= 5); + NVME_REQUIRE(stats.cache_hits >= 1); + NVME_REQUIRE(stats.inflight_deduplications + stats.cache_hits >= 2); + NVME_REQUIRE(stats.errors == 0); + NVME_REQUIRE(stats.payload_bytes > 0); + NVME_REQUIRE(stats.active_io_ns > 0); +} + +TEST_CASE(MoeNvmeSchedulerFixture, speculation_cannot_consume_demand_reserve) { + SyntheticModel model; + MoeNvmeConfig config; + config.backend = MoeNvmeBackend::Mmap; + config.direct_io = MoeNvmeDirectMode::Disabled; + config.host_slots = 4; + config.io_threads = 1; + config.demand_reserve = 2; + + MoeNvmeScheduler scheduler; + std::string err; + NVME_REQUIRE(scheduler.init(config, 32 * 1024, aligned_allocate, + aligned_free, nullptr, &err)); + NVME_REQUIRE(scheduler.bind_source( + {model.file.data(), model.file.size(), -1}, model.regions, &err)); + + NVME_REQUIRE(scheduler.request(0, 0, MoeNvmePriority::Prefetch, &err)); + NVME_REQUIRE(scheduler.request(0, 1, MoeNvmePriority::Prefetch, &err)); + NVME_REQUIRE(!scheduler.request(0, 2, MoeNvmePriority::Prefetch, nullptr)); + + // Demand always has admission rights and upgrades a speculative resident. + MoeNvmeLease demand; + NVME_REQUIRE(scheduler.acquire(0, 2, demand, &err)); + verify_lease(demand, 0, 2); + demand.reset(); + + const MoeNvmeStats stats = scheduler.stats(); + NVME_REQUIRE(stats.prefetch_drops >= 1); + NVME_REQUIRE(stats.errors == 0); +} + +TEST_CASE(MoeNvmeSchedulerFixture, split_model_reads_tensor_spans_from_multiple_shards) { + constexpr int experts = SyntheticModel::kExperts; + std::vector shard_a(512 * 1024, 0xa5); + std::vector shard_b(512 * 1024, 0x5a); + LayerExpertRegions layer; + layer.fused_gate_up = false; + layer.expert_bytes_gate = 4093; + layer.expert_bytes_up = 6141; + layer.expert_bytes_down = 8189; + layer.gate_exps = {257, layer.expert_bytes_gate * experts, 0}; + layer.up_exps = {129, layer.expert_bytes_up * experts, 1}; + layer.down_exps = { + layer.up_exps.offset + layer.up_exps.size + 197, + layer.expert_bytes_down * experts, 1}; + fill_tensor(shard_a, layer.gate_exps, layer.expert_bytes_gate, + 0, 0, experts); + fill_tensor(shard_b, layer.up_exps, layer.expert_bytes_up, + 0, 1, experts); + fill_tensor(shard_b, layer.down_exps, layer.expert_bytes_down, + 0, 2, experts); + + MoeExpertIoLayout layout; + std::string err; + NVME_REQUIRE(make_moe_expert_io_layout( + 0, 6, layer, std::vector{shard_a.size(), shard_b.size()}, + 4096, layout, &err)); + NVME_REQUIRE(layout.spans[0].source_index == 0); + NVME_REQUIRE(layout.spans[1].source_index == 1); + NVME_REQUIRE(layout.spans[2].source_index == 1); + + MoeNvmeConfig config; + config.backend = MoeNvmeBackend::Mmap; + config.direct_io = MoeNvmeDirectMode::Disabled; + config.host_slots = 4; + config.io_threads = 2; + MoeNvmeScheduler scheduler; + NVME_REQUIRE(scheduler.init(config, 32 * 1024, aligned_allocate, + aligned_free, nullptr, &err)); + const std::vector sources = { + {shard_a.data(), shard_a.size(), -1}, + {shard_b.data(), shard_b.size(), -1}, + }; + NVME_REQUIRE(scheduler.bind_sources(sources, {layer}, &err)); + MoeNvmeLease lease; + NVME_REQUIRE(scheduler.acquire(0, 6, lease, &err)); + verify_lease(lease, 0, 6); + lease.reset(); + NVME_REQUIRE(scheduler.stats().read_ops == 3); + NVME_REQUIRE(scheduler.stats().errors == 0); +} + +#if !defined(_WIN32) +TEST_CASE(MoeNvmeSchedulerFixture, split_real_files_use_the_fd_backend) { + constexpr int experts = SyntheticModel::kExperts; + std::vector shard_a(512 * 1024, 0xa5); + std::vector shard_b(512 * 1024, 0x5a); + LayerExpertRegions layer; + layer.expert_bytes_gate = 4093; + layer.expert_bytes_up = 6141; + layer.expert_bytes_down = 8189; + layer.gate_exps = {257, layer.expert_bytes_gate * experts, 0}; + layer.up_exps = {129, layer.expert_bytes_up * experts, 1}; + layer.down_exps = { + layer.up_exps.offset + layer.up_exps.size + 197, + layer.expert_bytes_down * experts, 1}; + fill_tensor(shard_a, layer.gate_exps, layer.expert_bytes_gate, + 0, 0, experts); + fill_tensor(shard_b, layer.up_exps, layer.expert_bytes_up, + 0, 1, experts); + fill_tensor(shard_b, layer.down_exps, layer.expert_bytes_down, + 0, 2, experts); + + char path_a[] = "/tmp/moe_nvme_shard_a_XXXXXX"; + char path_b[] = "/tmp/moe_nvme_shard_b_XXXXXX"; + const int fd_a = ::mkstemp(path_a); + const int fd_b = ::mkstemp(path_b); + NVME_REQUIRE(fd_a >= 0 && fd_b >= 0); + ::unlink(path_a); + ::unlink(path_b); + auto write_all = [](int fd, const std::vector & bytes) { + size_t written = 0; + while (written < bytes.size()) { + const ssize_t result = ::write( + fd, bytes.data() + written, bytes.size() - written); + NVME_REQUIRE(result > 0); + written += (size_t) result; + } + }; + write_all(fd_a, shard_a); + write_all(fd_b, shard_b); + + MoeNvmeConfig config; + config.backend = MoeNvmeBackend::Auto; + config.direct_io = MoeNvmeDirectMode::Disabled; + config.host_slots = 4; + config.io_threads = 2; + MoeNvmeScheduler scheduler; + std::string err; + NVME_REQUIRE(scheduler.init(config, 32 * 1024, aligned_allocate, + aligned_free, nullptr, &err)); + NVME_REQUIRE(scheduler.bind_sources({ + {shard_a.data(), shard_a.size(), fd_a}, + {shard_b.data(), shard_b.size(), fd_b}, + }, {layer}, &err)); + NVME_REQUIRE(std::string(scheduler.effective_backend_name()) != "mmap-workers"); + MoeNvmeLease lease; + NVME_REQUIRE(scheduler.acquire(0, 4, lease, &err)); + verify_lease(lease, 0, 4); + lease.reset(); + NVME_REQUIRE(scheduler.stats().read_ops == 3); + NVME_REQUIRE(scheduler.stats().errors == 0); + scheduler.destroy(); + ::close(fd_a); + ::close(fd_b); +} + +TEST_CASE(MoeNvmeSchedulerFixture, real_file_backend_reads_exact_bytes) { + SyntheticModel model; + char path[] = "/tmp/moe_nvme_scheduler_XXXXXX"; + const int fd = ::mkstemp(path); + NVME_REQUIRE(fd >= 0); + ::unlink(path); + size_t written = 0; + while (written < model.file.size()) { + const ssize_t result = ::write(fd, model.file.data() + written, + model.file.size() - written); + NVME_REQUIRE(result > 0); + written += (size_t) result; + } + + MoeNvmeConfig config; + config.backend = MoeNvmeBackend::Auto; + config.direct_io = MoeNvmeDirectMode::Disabled; + config.host_slots = 4; + config.io_threads = 2; + MoeNvmeScheduler scheduler; + std::string err; + NVME_REQUIRE(scheduler.init(config, 32 * 1024, aligned_allocate, + aligned_free, nullptr, &err)); + NVME_REQUIRE(scheduler.bind_source( + {model.file.data(), model.file.size(), fd}, model.regions, &err)); + MoeNvmeLease lease; + NVME_REQUIRE(scheduler.acquire(1, 7, lease, &err)); + verify_lease(lease, 1, 7); + lease.reset(); + NVME_REQUIRE(scheduler.stats().read_ops == 2); + NVME_REQUIRE(scheduler.stats().errors == 0); + scheduler.destroy(); + ::close(fd); +} +#endif From 2a11e8510c510dcc2a0c754d6bfb6bba90cd0811 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:15:15 +0200 Subject: [PATCH 02/20] perf(moe): generalize streamed experts for single-device Strix --- server/CMakeLists.txt | 22 + server/docs/DS4.md | 17 +- server/docs/ENVIRONMENT.md | 2 +- server/docs/KIMI_K3_HETERO.md | 28 +- server/docs/MOE_NVME_STREAMING.md | 99 ++- server/src/common/moe_hybrid_storage.h | 16 + server/src/common/moe_hybrid_stream.cpp | 786 ++++++++++++++++++++- server/src/common/moe_hybrid_stream.h | 85 +++ server/src/common/moe_nvme_scheduler.cpp | 86 ++- server/src/common/moe_nvme_scheduler.h | 26 + server/src/deepseek4/deepseek4_backend.cpp | 158 ++++- server/test/bench_kimi_k3_hetero.cpp | 169 ++--- server/test/test_moe_nvme_scheduler.cpp | 71 ++ server/test/test_moe_stream_compute.cpp | 311 ++++++++ 14 files changed, 1656 insertions(+), 220 deletions(-) create mode 100644 server/test/test_moe_stream_compute.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 6e1c0950d..fd1c6ffa2 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -719,6 +719,28 @@ if(DFLASH27B_TESTS) endif() endif() + # End-to-end numerical check for the reusable streamed-expert graph. It + # uses tiny generated F32 experts and defaults to GPU 0, so a Strix-only + # machine can run it without any model download. DFLASH_TEST_GPU selects a + # different device on multi-GPU qualification hosts. + if(UNIX AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_moe_stream_compute.cpp") + add_executable(test_moe_stream_compute + test/test_unit_main.cpp + test/test_moe_stream_compute.cpp) + target_include_directories(test_moe_stream_compute PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/src/common + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + target_link_libraries(test_moe_stream_compute PRIVATE + dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(test_moe_stream_compute PRIVATE CUDA::cudart) + else() + target_link_libraries(test_moe_stream_compute PRIVATE hip::host) + endif() + list(APPEND _raw_unit_test_targets test_moe_stream_compute) + endif() + # Kimi-K3 routed-core qualification. This replays the exact latent-expert # geometry (IQ1_S, 3584 -> 3072 -> 3584, SiTU) through the common NVMe # stream engine. It is deliberately independent of a Kimi model loader so diff --git a/server/docs/DS4.md b/server/docs/DS4.md index 388f89445..61fd5928b 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -177,12 +177,13 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. ### NVMe cold-capacity tier -When the cold expert stack cannot fit on Strix, the inference engine turns the -safe remaining Strix memory into an adaptive warm-expert cache and streams only -exact routed misses from NVMe. -`DFLASH_MOE_NVME_COLD_TIER=auto` selects this only when required by measured -free memory; `on` forces qualification and `off` requires resident experts. -The R9700 continues to own dense layers and hot experts. See +When the cold expert stack cannot fit on its compute device, the inference +engine turns safe remaining memory into an adaptive warm-expert cache and +streams only exact routed misses from NVMe. This supports both R9700+Strix +expert parallelism and a single Strix Halo. `DFLASH_MOE_NVME_COLD_TIER=auto` +selects streaming when capacity requires it; `on` forces at least one cold +expert per layer for qualification and `off` requires resident experts. On a +full Lucebox the R9700 continues to own dense layers and hot experts. See [`MOE_NVME_STREAMING.md`](MOE_NVME_STREAMING.md) for the data path, tuning, and benchmark methodology. @@ -197,8 +198,8 @@ and benchmark methodology. | `DFLASH_DS4_MOE_TP` | Enable routed-expert partitioning. | | `DFLASH_DS4_MOE_TP_INPROC` | Use two local HIP backends instead of an expert IPC worker. | | `DFLASH_DS4_MOE_TP_GPU` | HIP device that owns the cold expert stack. | -| `DFLASH_MOE_NVME_COLD_TIER` | `auto`, `on`, or `off` for the Strix-backed SSD cold-capacity tier. | -| `DFLASH_MOE_NVME_DEVICE_CACHE_MB` | Optional explicit Strix adaptive expert-cache budget; auto mode otherwise uses safe free memory. | +| `DFLASH_MOE_NVME_COLD_TIER` | `auto`, `on`, or `off` for the SSD cold-capacity tier on dual-device or Strix-only deployments. | +| `DFLASH_MOE_NVME_DEVICE_CACHE_MB` | Optional explicit adaptive device expert-cache budget; auto mode otherwise uses safe free memory. | | `DFLASH_EXPERT_BUDGET_MB` | Main-GPU memory budget for hot experts. | | `DFLASH_DS4_HOTNESS_CSV` | Optional per-layer routing profile for hot placement. | | `GGML_CUDA_BATCH_PEER_COPIES` | Batch ordered peer copies behind one dependency. | diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index e57ddb7e1..5cdafe579 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -28,7 +28,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_MMID_GROUPED_DEVICE` | -1 | Optional zero-based device restriction; unset/-1 applies to every eligible device. | | `DFLASH_DS4_MOE_TP` / `DFLASH_DS4_MOE_TP_INPROC` | unset | BURN-IN: enable DeepSeek4 route-owner expert parallelism in one process. | | `DFLASH_DS4_MOE_TP_GPU` | auto | HIP device for the cold DeepSeek4 expert owner. | -| `DFLASH_MOE_NVME_COLD_TIER` | auto | BURN-IN: DeepSeek4 cold-capacity policy (`auto`, `on`, `off`). Auto streams only when the cold stack does not fit on Strix with reserve. | +| `DFLASH_MOE_NVME_COLD_TIER` | auto | BURN-IN: DeepSeek4 cold-capacity policy (`auto`, `on`, `off`) for dual-device and Strix-only execution. | | `DFLASH_MOE_NVME_*` | tuned defaults | BURN-IN: bounded MoE SSD scheduler/backend controls; see `MOE_NVME_STREAMING.md`. | | `GGML_CUDA_BATCH_PEER_COPIES` | unset | BURN-IN: publish ordered HIP peer copies with one cross-device dependency per source/destination pair. | | `DFLASH_MOE_PREFILL_PERSISTENT_OWNER_ALLOC` | 1 for qualified long heterogeneous prefill | KILL SWITCH: =0 restores per-layer route/owner scratch allocation. | diff --git a/server/docs/KIMI_K3_HETERO.md b/server/docs/KIMI_K3_HETERO.md index ec4e45500..85ad15e3b 100644 --- a/server/docs/KIMI_K3_HETERO.md +++ b/server/docs/KIMI_K3_HETERO.md @@ -18,13 +18,16 @@ missing without making the scheduler Kimi-specific: - The common streamed graph supports SiTU as well as SwiGLU and DS4's clamped SwiGLU. - `bench_kimi_k3_hetero` runs Kimi's exact IQ1_S routed-expert geometry through - a reusable graph: 896 experts, top-16, 92 MoE layers, + the same model-neutral persistent evaluator used by production adapters: + 896 experts, top-16, 92 MoE layers, `3584 -> 3072 -> 3584`, and SiTU (`beta=4`, `linear_beta=25`). It overlaps SSD/H2D for expert N+1 with compute for expert N. -Six scheduler tests pass on the AMD Lucebox, including mmap and real-file -multi-shard reads. Existing single-file DS4 descriptors remain source index -zero and need no model-specific change. +Seven scheduler tests pass on the AMD Lucebox, including expert-major one-read +records, mmap, and real-file multi-shard reads. A separate numerical test +matches both tensor-major and expert-major streamed GPU execution against a +CPU oracle on gfx1151 and gfx1201. Existing single-file DS4 descriptors remain +source index zero and need no model-specific change. This is not yet a complete Kimi K3 backend. KDA/MLA, Attention Residuals, the vision encoder, the latent projections around the routed core, tokenizer, and @@ -67,16 +70,17 @@ the byte values do not affect transfer volume or kernel shape. | Owner of streamed experts | Scenario | Pipeline | Routed-core rate | |---|---|---:|---:| -| Strix Halo | 3 tokens, balanced cold routes, compute on | **3.716 GiB/s** | **0.420 token/s** | +| Strix Halo | 3 tokens, balanced cold routes, common evaluator | **3.804 GiB/s** | **0.430 token/s** | | R9700 | 1 token, balanced cold routes, compute on | 2.091 GiB/s | 0.236 token/s | | Strix Halo | 2 tokens, 10 GiB cache, unrelated routes | 3.638 GiB/s, 0.82% hits | 0.415 token/s | | Strix Halo | 2 tokens, 10 GiB cache, identical routes | 3.557 GiB/s, 50% aggregate hits | 0.804 token/s | -The sustained Strix run moved 26.531982 GiB in 7.139575 seconds, evaluated -4,416 routed experts, and reported zero I/O errors. Adding the exact expert -math did not reduce the cold result materially: compute is hidden behind the -8.84 GiB/token storage path. The R9700 is the wrong cold owner because the -extra discrete-GPU upload path cuts end-to-end throughput. +The sustained Strix run moved 26.531982 GiB in 6.974184 seconds, evaluated +4,416 routed experts with one graph build and 4,415 graph-cache hits, and +reported zero I/O errors. Adding the exact expert math did not reduce the cold +result materially: compute is hidden behind the 8.84 GiB/token storage path. +The R9700 is the wrong cold owner because the extra discrete-GPU upload path +cuts end-to-end throughput. The repeated-route result is deliberately a best case, not a prediction. Kimi K3 was designed for balanced expert use. With unrelated balanced routes, @@ -97,8 +101,8 @@ reserves. A simple placement is: After approximately 57.94 GiB of non-routed weights plus OS, workspace, and a moderate context reserve, roughly 70-85 GiB may remain for routed experts. Under a uniform balanced-routing assumption this covers about 14-17% of the -routed pool. At the measured 3.716 GiB/s, the storage-only ceiling is then -approximately 0.49-0.50 token/s. Full inference will be lower unless dense +routed pool. At the measured 3.804 GiB/s, the storage-only ceiling is then +approximately 0.50-0.51 token/s. Full inference will be lower unless dense R9700 work overlaps almost completely with Strix expert service. So the honest expectation for this quant is **roughly one token every two to diff --git a/server/docs/MOE_NVME_STREAMING.md b/server/docs/MOE_NVME_STREAMING.md index b0e1e8557..5b8ac960b 100644 --- a/server/docs/MOE_NVME_STREAMING.md +++ b/server/docs/MOE_NVME_STREAMING.md @@ -3,9 +3,9 @@ This is an inference-engine feature. It streams existing expert-weight bytes from SSD without changing their format or numerical representation. -## Lucebox data path +## Deployment shapes -The three owners have distinct jobs: +On a full Lucebox, the three owners have distinct jobs: 1. **R9700** owns dense layers and the statically hot routed experts. 2. **Strix Halo** owns an adaptive warm-expert cache, the bounded SSD staging @@ -15,12 +15,33 @@ The three owners have distinct jobs: 3. **NVMe** is the capacity tier for true cache misses that do not fit in the Strix safe-memory budget. -`LayerExpertRegions` is the model-adapter contract. It describes the exact -GGUF byte ranges for each layer's gate/up/down tensors (or fused gate+up). -The scheduler therefore has no model names, tensor names, expert dimensions, -or quantization-format assumptions. Each tensor range can also select a model -shard; single-file models use shard zero. See [KIMI_K3_HETERO.md](KIMI_K3_HETERO.md) -for the 14-shard Kimi K3 qualification. +On a Strix Halo-only machine, the same device owns dense/static-hot weights, +the warm cache, and streamed-expert execution. The planner reserves KV and +runtime headroom before assigning otherwise-unused UMA to the cache. Nothing +in the scheduler assumes an R9700 or peer access. + +## Model integration contract + +A model adapter supplies three independent descriptions: + +1. `LayerExpertRegions` describes physical bytes and model shards. Ordinary + GGUF uses tensor-major gate/up/down regions. An optional expert-major record + stores all components of one expert contiguously, reducing three reads to + one without changing weight values. Component offsets must respect the + target backend's tensor alignment. +2. `MoeStreamExpertSpec` describes dimensions, tensor types, scales, and the + gated activation. It supports separate or fused gate/up, different routed + input/output widths, SwiGLU, clamped SwiGLU, and SiTU. +3. `MoeStreamRouteBatch` carries the native router IDs, weights, and F32 input + activations. It contains no architecture-specific tensor names. + +`eval_moe_streamed_experts` validates the byte layout against the numerical +specification before compute. It then uses a bounded cache of persistent graphs +keyed by the full specification and active token width. New MoE families only +need an adapter that fills these descriptors; the storage scheduler, cache, +and compute pipeline remain unchanged. Each file range can select a different +shard, while single-file models use shard zero. See +[KIMI_K3_HETERO.md](KIMI_K3_HETERO.md) for the 14-shard Kimi K3 qualification. ## Scheduler @@ -45,10 +66,13 @@ for the 14-shard Kimi K3 qualification. upload and execution of expert N. - Two or more rotating GPU slots separate the expert being computed from the expert being uploaded. -- Otherwise-unused Strix memory becomes a contiguous model-neutral expert +- Otherwise-unused device memory becomes a contiguous model-neutral expert cache indexed by `(layer, expert)`. Cache hits issue neither SSD reads nor host-to-device copies. LFRU replacement cannot evict a pending upload or an expert currently executing. +- Persistent compute graphs remove graph construction and activation-buffer + allocation from the steady-state expert loop. A bounded LRU supports models + whose layers use more than one expert shape or quantization format. The native model router remains authoritative. Prediction may only issue a bounded prefetch; a wrong prediction cannot change model output. @@ -75,6 +99,25 @@ export DFLASH_MOE_NVME_COLD_TIER=auto Strix receives all currently usable memory (after reserve) as its adaptive expert-cache budget, and only the remainder spills to SSD. +### Strix Halo only + +Do not enable the MoE-TP variables. A single-device Strix machine normally +exposes its GPU as `hip:0`: + +```bash +unset DFLASH_DS4_MOE_TP DFLASH_DS4_MOE_TP_INPROC DFLASH_DS4_MOE_TP_GPU +export DFLASH_MOE_NVME_COLD_TIER=on + +./build-hip/dflash_server /path/to/model.gguf \ + --target-device hip:0 --max-ctx 8192 +``` + +`on` keeps at least one expert per layer in the SSD tier even if the model +would otherwise be fully resident, making the path directly testable. For a +model that genuinely exceeds UMA, `auto` chooses the partial placement itself. +`DFLASH_EXPERT_BUDGET_MB` can additionally cap static routed-weight residency; +`DFLASH_MOE_NVME_DEVICE_CACHE_MB` can override the adaptive cache budget. + ## Tuning and diagnostics Defaults are intentionally small: eight pinned host slots, four fallback I/O @@ -90,13 +133,19 @@ not improve the qualified P310 drive and consumes extra pinned/system memory. | `DFLASH_MOE_NVME_DEMAND_RESERVE` | `2` | Slots unavailable to speculation | | `DFLASH_MOE_NVME_PREFETCH_BATCH` | `2` | Maximum speculative jobs per ring submission | | `DFLASH_MOE_NVME_DEVICE_SLOTS` | `2` | Minimum rotating GPU expert buffers | -| `DFLASH_MOE_NVME_DEVICE_CACHE_MB` | automatic/`0` | Override adaptive Strix expert-cache memory; `0` leaves only pipeline slots | +| `DFLASH_MOE_NVME_DEVICE_CACHE_MB` | automatic/`0` | Override adaptive device expert-cache memory; `0` leaves only pipeline slots | +| `DFLASH_MOE_NVME_GRAPH_CACHE` | `8` | Persistent expert-graph variants retained per stream engine; `0` is a diagnostic no-cache mode | +| `DFLASH_MOE_NVME_REFERENCE_EVAL` | unset | Diagnostic only: `1` restores the allocation-heavy reference evaluator for numerical/performance A/B | Shutdown telemetry reports logical and physical bytes, measured read service -rate, cache hits, demand wait, de-duplication, dropped speculation, and errors. -The standalone targets `test_moe_nvme_scheduler`, `bench_moe_nvme_io`, and -`bench_moe_nvme_pipeline` test correctness, raw storage, and the complete -SSD-to-GPU path respectively. Benchmarks are read-only. +rate, cache hits, demand wait, de-duplication, dropped speculation, errors, and +persistent-graph builds/hits/evictions. `test_moe_stream_compute` generates +tiny experts and checks both tensor-major and expert-major GPU results against +a CPU oracle. It defaults to GPU 0, so it runs directly on Strix-only systems; +`DFLASH_TEST_GPU` selects another device on multi-GPU hosts. The standalone +targets `test_moe_nvme_scheduler`, `bench_moe_nvme_io`, and +`bench_moe_nvme_pipeline` test scheduling, raw storage, and the complete +SSD-to-GPU path. Benchmarks are read-only. ## Qualification result (2026-07-30) @@ -127,6 +176,16 @@ prompt issued no additional SSD reads and completed in 0.842 seconds: a 2.86x warm-request improvement with identical output. Final counters were 731 cache misses, 1,463 Strix hits, and zero I/O errors. +The single-Strix production path was also qualified with no MoE-TP variables, +a 12 GiB static-expert cap, a 255 MiB device cache, and 79.95 GiB left in the +SSD tier. Two identical seven-token prompts returned identical output in 3.762 +and 3.503 seconds of prefill. Across both requests the engine moved 25.367 GiB +at 4.350 GiB/s active I/O, reported zero errors, and used one graph build for +3,056 launches (3,055 graph-cache hits). The allocation-heavy reference path +took 3.816 and 3.599 seconds on the same run shape, so persistence improved +this deliberately cold, storage-bound case by 1.4-2.7%; its larger value is +removing thousands of allocations when more of the route set is warm. + ## Research lineage and next optimization The bounded priority/cache design follows the lessons of MoE-Infinity and @@ -136,9 +195,9 @@ that leaves the native router untouched. MoE-SpAc motivates compile-time expert layout and I/O coalescing. Tutti's slack-aware `io_uring` scheduling is relevant when persistent KV traffic shares the device. -The Kimi qualification now has a persistent reusable expert graph and proves -that its small expert math can hide under SSD service. The production DS4 path -still constructs a graph per streamed expert. Its next implementation step is -to move that qualification design into the common evaluator, then overlap -hot-owner work with cold-owner streaming. Any learned predictor comes after -that deterministic path is qualified. +The persistent graph and model-neutral evaluator are now shared by Kimi +qualification and production DS4 streaming. The next high-value work is a +repacked expert-major artifact, followed by overlap of hot-owner work with +cold-owner streaming. Any learned predictor comes after that deterministic +path is qualified, and may only prefetch routes selected later by the native +router. diff --git a/server/src/common/moe_hybrid_storage.h b/server/src/common/moe_hybrid_storage.h index e8558fb57..797ad26eb 100644 --- a/server/src/common/moe_hybrid_storage.h +++ b/server/src/common/moe_hybrid_storage.h @@ -24,6 +24,21 @@ struct ExpertFileRegion { uint32_t source_index = 0; }; +// Optional expert-major representation. All components of one expert occupy +// one contiguous record, allowing one storage request instead of two or three +// tensor-major requests. The numerical tensor types and values are unchanged; +// only their physical ordering differs. Model adapters can leave this disabled +// for ordinary GGUF tensor-major files. +struct ExpertMajorFileLayout { + ExpertFileRegion experts; + size_t expert_stride = 0; + size_t gate_offset = 0; + size_t up_offset = 0; + size_t down_offset = 0; + size_t gate_up_offset = 0; + bool enabled = false; +}; + // Per-layer file regions for all expert tensors (used by streaming prefill). struct LayerExpertRegions { ExpertFileRegion gate_exps; @@ -35,6 +50,7 @@ struct LayerExpertRegions { size_t expert_bytes_down = 0; size_t expert_bytes_gate_up = 0; bool fused_gate_up = false; + ExpertMajorFileLayout expert_major; }; // Cached FFN graph for a fixed number of selected experts. diff --git a/server/src/common/moe_hybrid_stream.cpp b/server/src/common/moe_hybrid_stream.cpp index f3bade16c..0bc86a9c1 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -6,10 +6,12 @@ #include "ggml-cuda.h" #include +#include #include #include #include #include +#include #include #include #include @@ -59,6 +61,36 @@ size_t align_up(size_t value, size_t alignment) { return (value + alignment - 1) & ~(alignment - 1); } +bool checked_mul_size(size_t a, size_t b, size_t & out) { + if (a != 0 && b > std::numeric_limits::max() / a) return false; + out = a * b; + return true; +} + +bool valid_ggml_type(ggml_type type) { + return type >= 0 && type < GGML_TYPE_COUNT; +} + +bool same_stream_spec(const MoeStreamExpertSpec & a, + const MoeStreamExpertSpec & b) { + return a.input_dim == b.input_dim && + a.intermediate_dim == b.intermediate_dim && + a.output_dim == b.output_dim && + a.gate_type == b.gate_type && + a.up_type == b.up_type && + a.down_type == b.down_type && + a.gate_up_type == b.gate_up_type && + a.fused_gate_up == b.fused_gate_up && + a.gated_activation == b.gated_activation && + a.swiglu_clamp == b.swiglu_clamp && + a.situ_beta == b.situ_beta && + a.situ_linear_beta == b.situ_linear_beta && + a.gate_scale == b.gate_scale && + a.up_scale == b.up_scale && + a.down_scale == b.down_scale && + a.gate_up_scale == b.gate_up_scale; +} + uint64_t device_key(int layer, int expert) { return ((uint64_t) (uint32_t) layer << 32) | (uint32_t) expert; } @@ -105,6 +137,8 @@ MoeStreamConfig MoeStreamConfig::from_env() { config.nvme = MoeNvmeConfig::from_env(config.nvme); config.device_slots = env_bounded_int( "DFLASH_MOE_NVME_DEVICE_SLOTS", config.device_slots, 2, 8); + config.graph_cache_entries = env_bounded_int( + "DFLASH_MOE_NVME_GRAPH_CACHE", config.graph_cache_entries, 0, 64); config.device_cache_bytes = env_mib( "DFLASH_MOE_NVME_DEVICE_CACHE_MB", config.device_cache_bytes); config.prefill_threshold = env_bounded_int( @@ -112,6 +146,383 @@ MoeStreamConfig MoeStreamConfig::from_env() { return config; } +bool make_moe_stream_expert_spec( + const MoeHybridConfig & cfg, + const MoeLayerDesc & desc, + const LayerExpertRegions & regions, + MoeStreamExpertSpec & out, + std::string * err) { + out = {}; + const int expert_dim = cfg.expert_embd(); + if (expert_dim <= 0 || cfg.n_ff_exp <= 0) { + if (err) *err = "streamed expert dimensions must be positive"; + return false; + } + out.input_dim = expert_dim; + out.intermediate_dim = cfg.n_ff_exp; + out.output_dim = expert_dim; + out.fused_gate_up = regions.fused_gate_up; + out.gated_activation = cfg.gated_activation; + out.swiglu_clamp = cfg.swiglu_clamp; + out.situ_beta = cfg.situ_beta; + out.situ_linear_beta = cfg.situ_linear_beta; + out.gate_scale = desc.ffn_gate_exps_s; + out.up_scale = desc.ffn_up_exps_s; + out.down_scale = desc.ffn_down_exps_s; + out.gate_up_scale = desc.ffn_gate_up_exps_s; + + if (out.fused_gate_up) { + if (!desc.ffn_gate_up_exps || !desc.ffn_down_exps) { + if (err) *err = "fused streamed expert is missing gate_up or down metadata"; + return false; + } + out.gate_up_type = desc.ffn_gate_up_exps->type; + out.down_type = desc.ffn_down_exps->type; + } else { + if (!desc.ffn_gate_exps || !desc.ffn_up_exps || !desc.ffn_down_exps) { + if (err) *err = "streamed expert is missing gate, up, or down metadata"; + return false; + } + out.gate_type = desc.ffn_gate_exps->type; + out.up_type = desc.ffn_up_exps->type; + out.down_type = desc.ffn_down_exps->type; + } + return true; +} + +bool validate_moe_stream_expert_layout( + const MoeStreamExpertSpec & spec, + const MoeExpertIoLayout & layout, + std::string * err) { + if (spec.input_dim <= 0 || spec.intermediate_dim <= 0 || + spec.output_dim <= 0 || !valid_ggml_type(spec.down_type) || + (spec.fused_gate_up + ? !valid_ggml_type(spec.gate_up_type) + : (!valid_ggml_type(spec.gate_type) || !valid_ggml_type(spec.up_type)))) { + if (err) *err = "invalid streamed expert shape or tensor type"; + return false; + } + if (spec.swiglu_clamp < 0.0f || + !std::isfinite(spec.swiglu_clamp) || + !std::isfinite(spec.gate_scale) || + !std::isfinite(spec.up_scale) || + !std::isfinite(spec.down_scale) || + !std::isfinite(spec.gate_up_scale) || + (spec.gated_activation != MoeGatedActivation::SwiGlu && + spec.gated_activation != MoeGatedActivation::Situ) || + (spec.gated_activation == MoeGatedActivation::Situ && + (spec.situ_beta <= 0.0f || spec.situ_linear_beta <= 0.0f || + !std::isfinite(spec.situ_beta) || + !std::isfinite(spec.situ_linear_beta)))) { + if (err) *err = "invalid streamed expert activation parameters"; + return false; + } + if (layout.fused_gate_up != spec.fused_gate_up) { + if (err) *err = "streamed storage and compute disagree about fused gate/up"; + return false; + } + + auto expected_bytes = [&](ggml_type type, int64_t columns, + int64_t rows, size_t & bytes) -> bool { + if (!valid_ggml_type(type) || columns <= 0 || rows <= 0) return false; + const size_t row = ggml_row_size(type, columns); + return checked_mul_size(row, (size_t) rows, bytes); + }; + auto require_component = [&](MoeExpertComponentKind kind, size_t expected, + const char * label) -> bool { + const MoeExpertComponentLayout * component = layout.component(kind); + if (!component || component->bytes != expected) { + if (err) { + *err = std::string("streamed ") + label + + " bytes do not match tensor type/shape"; + } + return false; + } + return true; + }; + + size_t down_bytes = 0; + if (!expected_bytes(spec.down_type, spec.intermediate_dim, + spec.output_dim, down_bytes)) { + if (err) *err = "streamed down tensor size overflow"; + return false; + } + if (spec.fused_gate_up) { + if (spec.intermediate_dim > std::numeric_limits::max() / 2) { + if (err) *err = "streamed fused intermediate dimension overflow"; + return false; + } + size_t gate_up_bytes = 0; + if (!expected_bytes(spec.gate_up_type, spec.input_dim, + 2LL * spec.intermediate_dim, gate_up_bytes) || + !require_component(MoeExpertComponentKind::FusedGateUp, + gate_up_bytes, "gate_up")) { + return false; + } + } else { + size_t gate_bytes = 0; + size_t up_bytes = 0; + if (!expected_bytes(spec.gate_type, spec.input_dim, + spec.intermediate_dim, gate_bytes) || + !expected_bytes(spec.up_type, spec.input_dim, + spec.intermediate_dim, up_bytes) || + !require_component(MoeExpertComponentKind::Gate, gate_bytes, "gate") || + !require_component(MoeExpertComponentKind::Up, up_bytes, "up")) { + return false; + } + } + return require_component(MoeExpertComponentKind::Down, down_bytes, "down"); +} + +namespace { + +class PersistentStreamExpertGraph { +public: + ~PersistentStreamExpertGraph() { destroy(); } + + bool matches(const MoeStreamExpertSpec & spec, int batch) const { + return batch_ == batch && same_stream_spec(spec_, spec); + } + + bool build(ggml_backend_t backend, + ggml_backend_buffer_t expert_buffer, + const MoeStreamExpertSpec & spec, + int batch, + const void * gate_data, + const void * up_data, + const void * down_data, + std::string * err) { + destroy(); + if (!backend || !expert_buffer || batch <= 0 || !gate_data || !down_data || + (!spec.fused_gate_up && !up_data)) { + if (err) *err = "invalid persistent streamed-expert graph arguments"; + return false; + } + backend_ = backend; + spec_ = spec; + batch_ = batch; + + ggml_init_params params{}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ctx_ = ggml_init(params); + if (!ctx_) { + if (err) *err = "ggml_init failed for persistent streamed expert"; + return false; + } + + input_ = ggml_new_tensor_2d( + ctx_, GGML_TYPE_F32, spec.input_dim, batch); + ggml_set_input(input_); + if (spec.fused_gate_up) { + gate_up_ = ggml_new_tensor_2d( + ctx_, spec.gate_up_type, spec.input_dim, + 2LL * spec.intermediate_dim); + down_ = ggml_new_tensor_2d( + ctx_, spec.down_type, spec.intermediate_dim, spec.output_dim); + ggml_set_input(gate_up_); + ggml_set_input(down_); + } else { + gate_ = ggml_new_tensor_2d( + ctx_, spec.gate_type, spec.input_dim, spec.intermediate_dim); + up_ = ggml_new_tensor_2d( + ctx_, spec.up_type, spec.input_dim, spec.intermediate_dim); + down_ = ggml_new_tensor_2d( + ctx_, spec.down_type, spec.intermediate_dim, spec.output_dim); + ggml_set_input(gate_); + ggml_set_input(up_); + ggml_set_input(down_); + } + + auto bind_external = [&](ggml_tensor * tensor, + const void * data, + const char * label) -> bool { + if (!tensor || !data) { + if (err) *err = std::string("invalid streamed ") + + label + " tensor binding"; + return false; + } + // Some GPU quant kernels require row-tail padding beyond + // ggml_nbytes(). The ordinary allocator supplies that padding, + // but a compact streamed record does not. Refuse such an adapter + // until it provides a padded device layout; otherwise a kernel + // could zero or read into the following component. + if (ggml_backend_buffer_get_alloc_size(expert_buffer, tensor) != + ggml_nbytes(tensor)) { + if (err) *err = std::string("streamed ") + label + + " tensor requires backend row padding; use a padded " + "device layout"; + return false; + } + const size_t alignment = + ggml_backend_buffer_get_alignment(expert_buffer); + if (alignment != 0 && + (uintptr_t) data % alignment != 0) { + if (err) *err = std::string("streamed ") + label + + " tensor is not aligned for the compute backend"; + return false; + } + if (ggml_backend_tensor_alloc( + expert_buffer, tensor, const_cast(data)) != + GGML_STATUS_SUCCESS) { + if (err) *err = std::string("failed to bind streamed ") + + label + " tensor to expert cache"; + return false; + } + return true; + }; + if (spec.fused_gate_up) { + if (!bind_external(gate_up_, gate_data, "gate_up") || + !bind_external(down_, down_data, "down")) { + return false; + } + } else if (!bind_external(gate_, gate_data, "gate") || + !bind_external(up_, up_data, "up") || + !bind_external(down_, down_data, "down")) { + return false; + } + + auto scale_if_needed = [&](ggml_tensor * value, float scale) { + return scale == 1.0f ? value : ggml_scale(ctx_, value, scale); + }; + auto gated_activation = [&](ggml_tensor * gate, + ggml_tensor * up) -> ggml_tensor * { + if (spec.gated_activation == MoeGatedActivation::Situ) { + ggml_tensor * nonlinear = ggml_scale( + ctx_, gate, 1.0f / spec.situ_beta); + nonlinear = ggml_tanh(ctx_, nonlinear); + nonlinear = ggml_scale(ctx_, nonlinear, spec.situ_beta); + nonlinear = ggml_mul( + ctx_, nonlinear, ggml_sigmoid(ctx_, gate)); + ggml_tensor * linear = ggml_scale( + ctx_, up, 1.0f / spec.situ_linear_beta); + linear = ggml_tanh(ctx_, linear); + linear = ggml_scale(ctx_, linear, spec.situ_linear_beta); + return ggml_mul(ctx_, nonlinear, linear); + } + if (spec.swiglu_clamp > 0.0f) { + return ggml_swiglu_ds4_split( + ctx_, gate, up, spec.swiglu_clamp); + } + return ggml_swiglu_split(ctx_, gate, up); + }; + + ggml_tensor * activated = nullptr; + if (gate_up_) { + ggml_tensor * combined = scale_if_needed( + ggml_mul_mat(ctx_, gate_up_, input_), spec.gate_up_scale); + ggml_tensor * gate_part = ggml_view_2d( + ctx_, combined, spec.intermediate_dim, batch, + combined->nb[1], 0); + ggml_tensor * up_part = ggml_view_2d( + ctx_, combined, spec.intermediate_dim, batch, + combined->nb[1], + (size_t) spec.intermediate_dim * sizeof(float)); + activated = gated_activation( + ggml_cont(ctx_, gate_part), ggml_cont(ctx_, up_part)); + } else { + ggml_tensor * gate_value = scale_if_needed( + ggml_mul_mat(ctx_, gate_, input_), spec.gate_scale); + ggml_tensor * up_value = scale_if_needed( + ggml_mul_mat(ctx_, up_, input_), spec.up_scale); + activated = gated_activation(gate_value, up_value); + } + output_ = scale_if_needed( + ggml_mul_mat(ctx_, down_, activated), spec.down_scale); + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_, 512, false); + ggml_build_forward_expand(graph_, output_); + alloc_ = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend_)); + if (!alloc_ || !ggml_gallocr_alloc_graph(alloc_, graph_)) { + if (err) *err = "persistent streamed-expert graph allocation failed"; + return false; + } + return true; + } + + bool launch(const void * gate_data, + const void * up_data, + const void * down_data, + const float * input, + std::string * err) { + if (!valid() || !gate_data || !down_data || !input || + (!spec_.fused_gate_up && !up_data)) { + if (err) *err = "persistent streamed-expert graph is not ready"; + return false; + } + if (gate_up_) gate_up_->data = const_cast(gate_data); + else gate_->data = const_cast(gate_data); + if (up_) up_->data = const_cast(up_data); + down_->data = const_cast(down_data); + size_t input_values = 0; + if (!checked_mul_size((size_t) spec_.input_dim, + (size_t) batch_, input_values)) { + if (err) *err = "streamed expert input size overflow"; + return false; + } + ggml_backend_tensor_set( + input_, input, 0, input_values * sizeof(float)); + if (ggml_backend_graph_compute_async(backend_, graph_) != + GGML_STATUS_SUCCESS) { + if (err) *err = "persistent streamed-expert graph launch failed"; + return false; + } + return true; + } + + bool finish(std::vector & output, std::string * err) { + if (!valid()) { + if (err) *err = "persistent streamed-expert graph is not ready"; + return false; + } + ggml_backend_synchronize(backend_); + size_t output_values = 0; + if (!checked_mul_size((size_t) spec_.output_dim, + (size_t) batch_, output_values)) { + if (err) *err = "streamed expert output size overflow"; + return false; + } + output.resize(output_values); + ggml_backend_tensor_get( + output_, output.data(), 0, output_values * sizeof(float)); + return true; + } + + void destroy() { + if (alloc_) ggml_gallocr_free(alloc_); + alloc_ = nullptr; + if (ctx_) ggml_free(ctx_); + ctx_ = nullptr; + graph_ = nullptr; + input_ = gate_ = up_ = down_ = gate_up_ = output_ = nullptr; + backend_ = nullptr; + batch_ = 0; + } + + bool valid() const { + return backend_ && ctx_ && graph_ && alloc_ && input_ && output_; + } + + uint64_t last_touch = 0; + +private: + ggml_backend_t backend_ = nullptr; + MoeStreamExpertSpec spec_{}; + int batch_ = 0; + ggml_context * ctx_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t alloc_ = nullptr; + ggml_tensor * input_ = nullptr; + ggml_tensor * gate_ = nullptr; + ggml_tensor * up_ = nullptr; + ggml_tensor * down_ = nullptr; + ggml_tensor * gate_up_ = nullptr; + ggml_tensor * output_ = nullptr; +}; + +} // namespace + struct MoeHybridStreamEngine::Runtime { struct DeviceSlot { void * data = nullptr; @@ -133,6 +544,7 @@ struct MoeHybridStreamEngine::Runtime { MoeStreamConfig config{}; std::unique_ptr io; cudaStream_t transfer_stream = nullptr; + ggml_backend_buffer_t device_pool_buffer = nullptr; void * device_pool = nullptr; size_t device_stride = 0; size_t device_pool_bytes = 0; @@ -143,6 +555,10 @@ struct MoeHybridStreamEngine::Runtime { uint64_t device_cache_misses = 0; uint64_t device_cache_evictions = 0; int active_slot = -1; + std::vector> graph_cache; + uint64_t graph_clock = 0; + MoeStreamComputeStats compute_stats{}; + std::mutex compute_mutex; }; template @@ -168,11 +584,16 @@ bool allocate_device_cache(RuntimeT & runtime, std::string * err) { // raced another allocation, converge to a smaller usable cache instead of // failing model startup. size_t attempt_slots = desired_slots; - cudaError_t gpu_err = cudaSuccess; + ggml_backend_buffer_type_t buft = + ggml_backend_get_default_buffer_type(runtime.backend); while (attempt_slots >= 2) { const size_t bytes = attempt_slots * runtime.device_stride; - gpu_err = cudaMalloc(&runtime.device_pool, bytes); - if (gpu_err == cudaSuccess) { + runtime.device_pool_buffer = ggml_backend_buft_alloc_buffer(buft, bytes); + if (runtime.device_pool_buffer) { + ggml_backend_buffer_set_usage( + runtime.device_pool_buffer, GGML_BACKEND_BUFFER_USAGE_COMPUTE); + runtime.device_pool = + ggml_backend_buffer_get_base(runtime.device_pool_buffer); runtime.device_pool_bytes = bytes; break; } @@ -180,17 +601,15 @@ bool allocate_device_cache(RuntimeT & runtime, std::string * err) { attempt_slots = std::max(2, attempt_slots * 3 / 4); } if (!runtime.device_pool) { - if (err) { - *err = std::string("failed to allocate SSD GPU expert cache: ") + - cudaGetErrorString(gpu_err); - } + if (err) *err = "failed to allocate SSD GPU expert cache"; return false; } try { runtime.device_slots.resize(attempt_slots); } catch (const std::bad_alloc &) { - (void) cudaFree(runtime.device_pool); + ggml_backend_buffer_free(runtime.device_pool_buffer); + runtime.device_pool_buffer = nullptr; runtime.device_pool = nullptr; runtime.device_pool_bytes = 0; if (err) *err = "failed to allocate SSD GPU cache metadata"; @@ -351,9 +770,11 @@ void MoeHybridStreamEngine::destroy() { const size_t device_cache_slot_count = runtime_->device_slots.size(); const size_t device_cache_byte_count = runtime_->device_pool_bytes; ScopedGpuDevice device_scope(runtime_->device); + if (runtime_->backend) ggml_backend_synchronize(runtime_->backend); if (runtime_->transfer_stream) { (void) cudaStreamSynchronize(runtime_->transfer_stream); } + runtime_->graph_cache.clear(); for (Runtime::DeviceSlot & slot : runtime_->device_slots) { slot.host_lease.reset(); if (slot.ready) (void) cudaEventDestroy(slot.ready); @@ -363,7 +784,10 @@ void MoeHybridStreamEngine::destroy() { } runtime_->device_slots.clear(); runtime_->device_index.clear(); - if (runtime_->device_pool) (void) cudaFree(runtime_->device_pool); + if (runtime_->device_pool_buffer) { + ggml_backend_buffer_free(runtime_->device_pool_buffer); + } + runtime_->device_pool_buffer = nullptr; runtime_->device_pool = nullptr; if (runtime_->transfer_stream) (void) cudaStreamDestroy(runtime_->transfer_stream); runtime_->transfer_stream = nullptr; @@ -387,7 +811,8 @@ void MoeHybridStreamEngine::destroy() { "payload=%.3f GiB physical=%.3f GiB active-io-rate=%.3f GiB/s " "cache-hit=%.1f%% mean-demand-wait=%.3f ms " "dedupe=%llu upgrades=%llu dropped-prefetch=%llu errors=%llu " - "strix-cache=%.1f MiB slots=%zu hits=%llu misses=%llu evictions=%llu\n", + "device-cache=%.1f MiB slots=%zu hits=%llu misses=%llu evictions=%llu " + "graphs=%llu graph-hits=%llu graph-evictions=%llu launches=%llu\n", runtime_->io->effective_backend_name(), (unsigned long long) stats.requests, (unsigned long long) stats.read_ops, @@ -400,7 +825,11 @@ void MoeHybridStreamEngine::destroy() { device_cache_slot_count, (unsigned long long) runtime_->device_cache_hits, (unsigned long long) runtime_->device_cache_misses, - (unsigned long long) runtime_->device_cache_evictions); + (unsigned long long) runtime_->device_cache_evictions, + (unsigned long long) runtime_->compute_stats.graph_builds, + (unsigned long long) runtime_->compute_stats.graph_cache_hits, + (unsigned long long) runtime_->compute_stats.graph_evictions, + (unsigned long long) runtime_->compute_stats.graph_launches); } runtime_->io->destroy(); } @@ -642,7 +1071,7 @@ bool MoeHybridStreamEngine::activate_device_slot(int device_slot, slot.pending = false; slot.host_lease.reset(); } - if (slot.layout.span_count < 2) { + if (slot.layout.component_count < 2) { if (err) *err = "SSD device slot has no complete expert"; return false; } @@ -706,40 +1135,62 @@ bool MoeHybridStreamEngine::stream_expert_sync( const void * MoeHybridStreamEngine::scratch_gate_data() const { if (!runtime_ || runtime_->active_slot < 0) return nullptr; const Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) runtime_->active_slot]; - return static_cast(slot.data) + slot.layout.spans[0].device_offset; + const MoeExpertComponentKind kind = slot.layout.fused_gate_up + ? MoeExpertComponentKind::FusedGateUp : MoeExpertComponentKind::Gate; + const MoeExpertComponentLayout * component = slot.layout.component(kind); + return component + ? static_cast(slot.data) + component->device_offset + : nullptr; } const void * MoeHybridStreamEngine::scratch_up_data() const { if (!runtime_ || runtime_->active_slot < 0) return nullptr; const Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) runtime_->active_slot]; if (slot.layout.fused_gate_up) return nullptr; - return static_cast(slot.data) + slot.layout.spans[1].device_offset; + const MoeExpertComponentLayout * component = + slot.layout.component(MoeExpertComponentKind::Up); + return component + ? static_cast(slot.data) + component->device_offset + : nullptr; } const void * MoeHybridStreamEngine::scratch_down_data() const { if (!runtime_ || runtime_->active_slot < 0) return nullptr; const Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) runtime_->active_slot]; - const int index = slot.layout.fused_gate_up ? 1 : 2; - return static_cast(slot.data) + slot.layout.spans[index].device_offset; + const MoeExpertComponentLayout * component = + slot.layout.component(MoeExpertComponentKind::Down); + return component + ? static_cast(slot.data) + component->device_offset + : nullptr; } size_t MoeHybridStreamEngine::scratch_gate_bytes() const { if (!runtime_ || runtime_->active_slot < 0) return 0; - return runtime_->device_slots[(size_t) runtime_->active_slot].layout.spans[0].bytes; + const MoeExpertIoLayout & layout = + runtime_->device_slots[(size_t) runtime_->active_slot].layout; + const MoeExpertComponentKind kind = layout.fused_gate_up + ? MoeExpertComponentKind::FusedGateUp : MoeExpertComponentKind::Gate; + const MoeExpertComponentLayout * component = layout.component(kind); + return component ? component->bytes : 0; } size_t MoeHybridStreamEngine::scratch_up_bytes() const { if (!runtime_ || runtime_->active_slot < 0) return 0; const MoeExpertIoLayout & layout = runtime_->device_slots[(size_t) runtime_->active_slot].layout; - return layout.fused_gate_up ? 0 : layout.spans[1].bytes; + if (layout.fused_gate_up) return 0; + const MoeExpertComponentLayout * component = + layout.component(MoeExpertComponentKind::Up); + return component ? component->bytes : 0; } size_t MoeHybridStreamEngine::scratch_down_bytes() const { if (!runtime_ || runtime_->active_slot < 0) return 0; const MoeExpertIoLayout & layout = runtime_->device_slots[(size_t) runtime_->active_slot].layout; - return layout.spans[layout.fused_gate_up ? 1 : 2].bytes; + const MoeExpertComponentLayout * component = + layout.component(MoeExpertComponentKind::Down); + return component ? component->bytes : 0; } size_t MoeHybridStreamEngine::pinned_bytes() const { @@ -758,7 +1209,11 @@ MoeNvmeStats MoeHybridStreamEngine::io_stats() const { return runtime_ && runtime_->io ? runtime_->io->stats() : MoeNvmeStats{}; } -bool eval_moe_cold_experts_streaming( +MoeStreamComputeStats MoeHybridStreamEngine::compute_stats() const { + return runtime_ ? runtime_->compute_stats : MoeStreamComputeStats{}; +} + +static bool eval_moe_cold_experts_streaming_reference( MoeHybridStreamEngine & engine, ggml_backend_t gpu_backend, const void * mmap_data, @@ -1021,4 +1476,295 @@ bool eval_moe_cold_experts_streaming( return true; } +bool eval_moe_streamed_experts( + MoeHybridStreamEngine & engine, + const MoeStreamExpertSpec & spec, + const MoeStreamRouteBatch & batch, + std::vector & out, + std::string * err) { + if (!engine.runtime_ || !engine.is_bound()) { + if (err) *err = "streamed expert evaluation requires a bound model source"; + return false; + } + if (batch.layer < 0 || batch.n_expert <= 0 || batch.top_k <= 0 || + batch.top_k > batch.n_expert || + batch.n_tokens <= 0 || !batch.inputs || !batch.selected_ids || + !batch.selected_weights || spec.input_dim <= 0 || + spec.output_dim <= 0) { + if (err) *err = "invalid model-neutral streamed route batch"; + return false; + } + if (batch.resident_local_by_global && + batch.resident_map_size < (size_t) batch.n_expert) { + if (err) *err = "streamed route residency map is smaller than n_expert"; + return false; + } + size_t output_values = 0; + if (!checked_mul_size((size_t) spec.output_dim, + (size_t) batch.n_tokens, output_values)) { + if (err) *err = "streamed route output size overflow"; + return false; + } + out.assign(output_values, 0.0f); + + auto & runtime = *engine.runtime_; + std::lock_guard compute_guard(runtime.compute_mutex); + + size_t route_slots = 0; + if (!checked_mul_size((size_t) batch.top_k, + (size_t) batch.n_tokens, route_slots)) { + if (err) *err = "streamed route slot count overflow"; + return false; + } + std::vector needed((size_t) batch.n_expert, false); + for (size_t i = 0; i < route_slots; ++i) { + const int32_t expert = batch.selected_ids[i]; + if (expert < 0) continue; + if (expert >= batch.n_expert) { + if (err) *err = "native router selected an out-of-range expert"; + return false; + } + if (!std::isfinite(batch.selected_weights[i])) { + if (err) *err = "native router produced a non-finite expert weight"; + return false; + } + if (batch.selected_weights[i] == 0.0f) continue; + if (batch.resident_local_by_global && + batch.resident_local_by_global[(size_t) expert] >= 0) { + continue; + } + needed[(size_t) expert] = true; + } + + std::vector unique_experts; + for (int expert = 0; expert < batch.n_expert; ++expert) { + if (needed[(size_t) expert]) unique_experts.push_back((int32_t) expert); + } + if (unique_experts.empty()) return true; + + engine.request_experts(batch.layer, unique_experts.data(), + (int) unique_experts.size(), + MoeNvmePriority::Demand); + int staged_slot = -1; + if (!engine.stage_expert_cached_async( + batch.layer, unique_experts[0], &staged_slot, err)) { + return false; + } + + auto acquire_graph = [&](int graph_batch, + std::unique_ptr & ephemeral, + PersistentStreamExpertGraph ** graph_out) -> bool { + *graph_out = nullptr; + const uint64_t touch = ++runtime.graph_clock; + if (runtime.config.graph_cache_entries > 0) { + for (auto & candidate : runtime.graph_cache) { + if (candidate && candidate->matches(spec, graph_batch)) { + candidate->last_touch = touch; + ++runtime.compute_stats.graph_cache_hits; + *graph_out = candidate.get(); + return true; + } + } + } + + const int active = runtime.active_slot; + if (active < 0 || active >= (int) runtime.device_slots.size()) { + if (err) *err = "no active streamed expert slot for graph build"; + return false; + } + const MoeExpertIoLayout & layout = + runtime.device_slots[(size_t) active].layout; + if (!validate_moe_stream_expert_layout(spec, layout, err)) return false; + + std::unique_ptr built( + new (std::nothrow) PersistentStreamExpertGraph); + if (!built) { + if (err) *err = "failed to allocate persistent streamed-expert graph"; + return false; + } + if (!built->build(runtime.backend, runtime.device_pool_buffer, + spec, graph_batch, + engine.scratch_gate_data(), + engine.scratch_up_data(), + engine.scratch_down_data(), err)) { + return false; + } + built->last_touch = touch; + ++runtime.compute_stats.graph_builds; + + if (runtime.config.graph_cache_entries <= 0) { + *graph_out = built.get(); + ephemeral = std::move(built); + return true; + } + if ((int) runtime.graph_cache.size() >= + runtime.config.graph_cache_entries) { + auto victim = std::min_element( + runtime.graph_cache.begin(), runtime.graph_cache.end(), + [](const auto & a, const auto & b) { + return a->last_touch < b->last_touch; + }); + if (victim != runtime.graph_cache.end()) { + runtime.graph_cache.erase(victim); + ++runtime.compute_stats.graph_evictions; + } + } + *graph_out = built.get(); + runtime.graph_cache.push_back(std::move(built)); + return true; + }; + + for (size_t expert_index = 0; + expert_index < unique_experts.size(); ++expert_index) { + const int current_slot = staged_slot; + if (!engine.activate_device_slot(current_slot, err)) return false; + auto release_current = [&]() { + engine.release_device_slot(current_slot); + }; + + struct TokenHit { int token; float weight; }; + std::vector hits; + hits.reserve((size_t) batch.n_tokens); + const int32_t expert = unique_experts[expert_index]; + for (int token = 0; token < batch.n_tokens; ++token) { + float combined_weight = 0.0f; + for (int rank = 0; rank < batch.top_k; ++rank) { + const size_t route = + (size_t) token * (size_t) batch.top_k + (size_t) rank; + if (batch.selected_ids[route] != expert) continue; + combined_weight += batch.selected_weights[route]; + } + if (!std::isfinite(combined_weight)) { + if (err) *err = "combined expert route weight overflowed"; + release_current(); + return false; + } + if (combined_weight != 0.0f) { + hits.push_back({token, combined_weight}); + } + } + if (hits.empty()) { + release_current(); + continue; + } + + size_t input_values = 0; + if (!checked_mul_size((size_t) spec.input_dim, hits.size(), + input_values)) { + if (err) *err = "streamed compact input size overflow"; + release_current(); + return false; + } + std::vector compact_input(input_values); + for (size_t i = 0; i < hits.size(); ++i) { + const float * src = batch.inputs + + (size_t) hits[i].token * (size_t) spec.input_dim; + std::memcpy(compact_input.data() + i * (size_t) spec.input_dim, + src, sizeof(float) * (size_t) spec.input_dim); + } + + std::unique_ptr ephemeral; + PersistentStreamExpertGraph * graph = nullptr; + if (!acquire_graph((int) hits.size(), ephemeral, &graph)) { + release_current(); + return false; + } + if (!validate_moe_stream_expert_layout( + spec, runtime.device_slots[(size_t) current_slot].layout, err) || + !graph->launch(engine.scratch_gate_data(), + engine.scratch_up_data(), + engine.scratch_down_data(), + compact_input.data(), err)) { + release_current(); + return false; + } + ++runtime.compute_stats.graph_launches; + + // Compute N is running while the already-admitted read for N+1 is + // acquired and uploaded into a different, eviction-protected slot. + if (expert_index + 1 < unique_experts.size()) { + int next_slot = -1; + if (!engine.stage_expert_cached_async( + batch.layer, unique_experts[expert_index + 1], + &next_slot, err)) { + ggml_backend_synchronize(runtime.backend); + release_current(); + return false; + } + staged_slot = next_slot; + } + + std::vector result; + if (!graph->finish(result, err)) { + release_current(); + return false; + } + for (size_t i = 0; i < hits.size(); ++i) { + float * dst = out.data() + + (size_t) hits[i].token * (size_t) spec.output_dim; + const float * src = result.data() + + i * (size_t) spec.output_dim; + const float weight = hits[i].weight; + for (int j = 0; j < spec.output_dim; ++j) { + dst[j] += weight * src[(size_t) j]; + } + } + release_current(); + } + return true; +} + +bool eval_moe_cold_experts_streaming( + MoeHybridStreamEngine & engine, + ggml_backend_t gpu_backend, + const void * mmap_data, + size_t mmap_size, + const MoeHybridConfig & cfg, + const MoeLayerDesc & desc, + const LayerExpertRegions & regions, + const MoeHybridLayerStorage & storage, + const float * cur_host, + const int32_t * selected_ids, + const float * selected_weights, + int n_tokens, + std::vector & out, + std::string * err, + int layer) { + const char * reference = std::getenv("DFLASH_MOE_NVME_REFERENCE_EVAL"); + if (reference && (std::strcmp(reference, "1") == 0 || + std::strcmp(reference, "on") == 0 || + std::strcmp(reference, "true") == 0)) { + return eval_moe_cold_experts_streaming_reference( + engine, gpu_backend, mmap_data, mmap_size, cfg, desc, regions, + storage, cur_host, selected_ids, selected_weights, n_tokens, + out, err, layer); + } + + int bound_layer = layer; + if (!engine.is_bound()) { + if (!mmap_data || mmap_size == 0 || + !engine.bind_sources({{mmap_data, mmap_size, -1}}, {regions}, err)) { + if (err && err->empty()) *err = "stream engine has no model source"; + return false; + } + bound_layer = 0; + } + + MoeStreamExpertSpec spec; + if (!make_moe_stream_expert_spec(cfg, desc, regions, spec, err)) return false; + MoeStreamRouteBatch route_batch; + route_batch.layer = bound_layer; + route_batch.n_expert = cfg.n_expert; + route_batch.top_k = cfg.n_expert_used; + route_batch.n_tokens = n_tokens; + route_batch.inputs = cur_host; + route_batch.selected_ids = selected_ids; + route_batch.selected_weights = selected_weights; + route_batch.resident_local_by_global = + storage.hot_local_by_global.empty() + ? nullptr : storage.hot_local_by_global.data(); + route_batch.resident_map_size = storage.hot_local_by_global.size(); + return eval_moe_streamed_experts(engine, spec, route_batch, out, err); +} + } // namespace dflash::common diff --git a/server/src/common/moe_hybrid_stream.h b/server/src/common/moe_hybrid_stream.h index 0d1c1f73c..e267c6867 100644 --- a/server/src/common/moe_hybrid_stream.h +++ b/server/src/common/moe_hybrid_stream.h @@ -25,6 +25,10 @@ struct MoeStreamConfig { int prefill_threshold = 8; int prefetch_layers = 2; int device_slots = 2; // double buffering is the minimum useful pipeline + // Persistent compute graphs are keyed by tensor types, dimensions, + // activation, scales, and batch width. A small bounded cache removes graph + // construction from decode without assuming every layer uses one format. + int graph_cache_entries = 8; // Optional adaptive GPU expert-cache budget. Zero keeps only the pipeline // slots. The hardware planner can safely assign otherwise-unused Strix // memory here while retaining its KV/graph reserve. @@ -34,6 +38,70 @@ struct MoeStreamConfig { static MoeStreamConfig from_env(); }; +// Complete numerical contract for one streamed gated expert. This is the +// model-adapter boundary: storage scheduling consumes byte ranges, while the +// reusable compute path consumes this shape/type/activation description. +// input_dim and output_dim may differ, although common routed FFNs use the +// same value for both. +struct MoeStreamExpertSpec { + int input_dim = 0; + int intermediate_dim = 0; + int output_dim = 0; + ggml_type gate_type = GGML_TYPE_COUNT; + ggml_type up_type = GGML_TYPE_COUNT; + ggml_type down_type = GGML_TYPE_COUNT; + ggml_type gate_up_type = GGML_TYPE_COUNT; + bool fused_gate_up = false; + MoeGatedActivation gated_activation = MoeGatedActivation::SwiGlu; + float swiglu_clamp = 0.0f; + float situ_beta = 4.0f; + float situ_linear_beta = 25.0f; + float gate_scale = 1.0f; + float up_scale = 1.0f; + float down_scale = 1.0f; + float gate_up_scale = 1.0f; +}; + +// Build the common contract from an existing model's ordinary MoE metadata. +// New adapters may either use this helper or populate MoeStreamExpertSpec +// directly when their latent input/output dimensions differ. +bool make_moe_stream_expert_spec( + const MoeHybridConfig & cfg, + const MoeLayerDesc & desc, + const LayerExpertRegions & regions, + MoeStreamExpertSpec & out, + std::string * err = nullptr); + +// Reject a shape/type/layout mismatch before launching a kernel. This makes a +// bad adapter fail deterministically instead of silently interpreting the +// wrong number of weight bytes. +bool validate_moe_stream_expert_layout( + const MoeStreamExpertSpec & spec, + const MoeExpertIoLayout & layout, + std::string * err = nullptr); + +// Model-neutral routed batch. Inputs and outputs are token-major contiguous +// F32 matrices. resident_local_by_global is optional; when supplied, IDs with +// a non-negative entry are already owned by another tier and are skipped. +struct MoeStreamRouteBatch { + int layer = 0; + int n_expert = 0; + int top_k = 0; + int n_tokens = 0; + const float * inputs = nullptr; + const int32_t * selected_ids = nullptr; + const float * selected_weights = nullptr; + const int32_t * resident_local_by_global = nullptr; + size_t resident_map_size = 0; +}; + +struct MoeStreamComputeStats { + uint64_t graph_builds = 0; + uint64_t graph_cache_hits = 0; + uint64_t graph_evictions = 0; + uint64_t graph_launches = 0; +}; + class MoeHybridStreamEngine { public: MoeHybridStreamEngine(); @@ -114,12 +182,29 @@ class MoeHybridStreamEngine { size_t scratch_bytes() const; const char * io_backend_name() const; MoeNvmeStats io_stats() const; + MoeStreamComputeStats compute_stats() const; private: + friend bool eval_moe_streamed_experts( + MoeHybridStreamEngine &, + const MoeStreamExpertSpec &, + const MoeStreamRouteBatch &, + std::vector &, + std::string *); struct Runtime; std::unique_ptr runtime_; }; +// Evaluate exactly the experts selected by the native router. Placement only +// decides which selected IDs arrive here; no prediction or cache policy can +// change the returned mathematical function. +bool eval_moe_streamed_experts( + MoeHybridStreamEngine & engine, + const MoeStreamExpertSpec & spec, + const MoeStreamRouteBatch & batch, + std::vector & out, + std::string * err = nullptr); + // Evaluate the cold contribution for one layer. All routed SSD requests are // admitted before compute starts, then double-buffered H2D runs concurrently // with the preceding expert graph. diff --git a/server/src/common/moe_nvme_scheduler.cpp b/server/src/common/moe_nvme_scheduler.cpp index bfc481b9a..540781d66 100644 --- a/server/src/common/moe_nvme_scheduler.cpp +++ b/server/src/common/moe_nvme_scheduler.cpp @@ -431,7 +431,7 @@ bool make_moe_expert_io_layout( MoeExpertIoLayout & out, std::string * err) { - out = {}; + out = MoeExpertIoLayout{}; out.key = { (int32_t) layer, (int32_t) expert }; out.fused_gate_up = regions.fused_gate_up; if (layer < 0 || expert < 0) { @@ -501,13 +501,85 @@ bool make_moe_expert_io_layout( return true; }; - if (regions.fused_gate_up) { - if (!add_span(regions.gate_up_exps, regions.expert_bytes_gate_up, "gate_up")) return false; + auto add_component = [&](MoeExpertComponentKind kind, size_t offset, + size_t bytes, size_t record_bytes, + const char * label) -> bool { + if (out.component_count >= 3 || bytes == 0 || + !range_in_bounds(offset, bytes, record_bytes)) { + if (err) *err = std::string("invalid ") + label + + " component in expert record"; + return false; + } + for (int i = 0; i < out.component_count; ++i) { + const MoeExpertComponentLayout & prior = out.components[i]; + const size_t prior_end = prior.device_offset + prior.bytes; + const size_t end = offset + bytes; + if (offset < prior_end && prior.device_offset < end) { + if (err) *err = std::string(label) + + " overlaps another expert component"; + return false; + } + } + out.components[out.component_count++] = {kind, offset, bytes}; + return true; + }; + + if (regions.expert_major.enabled) { + const ExpertMajorFileLayout & packed = regions.expert_major; + if (!add_span(packed.experts, packed.expert_stride, "expert-major")) { + return false; + } + if (regions.fused_gate_up) { + if (!add_component(MoeExpertComponentKind::FusedGateUp, + packed.gate_up_offset, + regions.expert_bytes_gate_up, + packed.expert_stride, "gate_up")) { + return false; + } + } else { + if (!add_component(MoeExpertComponentKind::Gate, + packed.gate_offset, regions.expert_bytes_gate, + packed.expert_stride, "gate") || + !add_component(MoeExpertComponentKind::Up, + packed.up_offset, regions.expert_bytes_up, + packed.expert_stride, "up")) { + return false; + } + } + if (!add_component(MoeExpertComponentKind::Down, + packed.down_offset, regions.expert_bytes_down, + packed.expert_stride, "down")) { + return false; + } + } else if (regions.fused_gate_up) { + if (!add_span(regions.gate_up_exps, regions.expert_bytes_gate_up, "gate_up") || + !add_component(MoeExpertComponentKind::FusedGateUp, + out.spans[out.span_count - 1].device_offset, + regions.expert_bytes_gate_up, device_cursor, + "gate_up")) { + return false; + } } else { - if (!add_span(regions.gate_exps, regions.expert_bytes_gate, "gate")) return false; - if (!add_span(regions.up_exps, regions.expert_bytes_up, "up")) return false; + if (!add_span(regions.gate_exps, regions.expert_bytes_gate, "gate") || + !add_component(MoeExpertComponentKind::Gate, + out.spans[out.span_count - 1].device_offset, + regions.expert_bytes_gate, device_cursor, "gate") || + !add_span(regions.up_exps, regions.expert_bytes_up, "up") || + !add_component(MoeExpertComponentKind::Up, + out.spans[out.span_count - 1].device_offset, + regions.expert_bytes_up, device_cursor, "up")) { + return false; + } + } + + if (!regions.expert_major.enabled) { + if (!add_span(regions.down_exps, regions.expert_bytes_down, "down") || + !add_component(MoeExpertComponentKind::Down, + out.spans[out.span_count - 1].device_offset, + regions.expert_bytes_down, device_cursor, "down")) { + return false; + } } - if (!add_span(regions.down_exps, regions.expert_bytes_down, "down")) return false; out.payload_bytes = device_cursor; out.host_bytes = host_cursor; @@ -1093,7 +1165,7 @@ void MoeNvmeLease::reset() { data_ = nullptr; slot_ = -1; generation_ = 0; - layout_ = {}; + layout_ = MoeExpertIoLayout{}; } MoeNvmeScheduler::MoeNvmeScheduler() : impl_(new Impl) {} diff --git a/server/src/common/moe_nvme_scheduler.h b/server/src/common/moe_nvme_scheduler.h index 4e3537a36..f803c3c1d 100644 --- a/server/src/common/moe_nvme_scheduler.h +++ b/server/src/common/moe_nvme_scheduler.h @@ -89,13 +89,39 @@ struct MoeExpertIoSpan { size_t io_bytes = 0; // aligned direct-I/O length }; +enum class MoeExpertComponentKind : uint8_t { + Gate, + Up, + Down, + FusedGateUp, +}; + +// Compute-facing view of one tensor inside the packed device slot. Keeping +// components separate from I/O spans lets tensor-major storage use three reads +// while an expert-major representation uses one read with identical compute +// pointers and numerical behavior. +struct MoeExpertComponentLayout { + MoeExpertComponentKind kind = MoeExpertComponentKind::Gate; + size_t device_offset = 0; + size_t bytes = 0; +}; + struct MoeExpertIoLayout { MoeExpertKey key; MoeExpertIoSpan spans[3]{}; int span_count = 0; + MoeExpertComponentLayout components[3]{}; + int component_count = 0; size_t payload_bytes = 0; size_t host_bytes = 0; bool fused_gate_up = false; + + const MoeExpertComponentLayout * component(MoeExpertComponentKind kind) const { + for (int i = 0; i < component_count; ++i) { + if (components[i].kind == kind) return &components[i]; + } + return nullptr; + } }; // Convert the common LayerExpertRegions descriptor into an exact read plan. diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index e530a438e..0c00ff228 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -40,6 +40,29 @@ static bool env_flag_enabled(const char * name) { return value && value[0] && std::strcmp(value, "0") != 0; } +enum class MoeNvmeColdTierMode { + Auto, + Enabled, + Disabled, + Invalid, +}; + +static MoeNvmeColdTierMode moe_nvme_cold_tier_mode() { + const char * value = std::getenv("DFLASH_MOE_NVME_COLD_TIER"); + if (!value || !value[0] || std::strcmp(value, "auto") == 0) { + return MoeNvmeColdTierMode::Auto; + } + if (std::strcmp(value, "1") == 0 || std::strcmp(value, "on") == 0 || + std::strcmp(value, "true") == 0) { + return MoeNvmeColdTierMode::Enabled; + } + if (std::strcmp(value, "0") == 0 || std::strcmp(value, "off") == 0 || + std::strcmp(value, "false") == 0) { + return MoeNvmeColdTierMode::Disabled; + } + return MoeNvmeColdTierMode::Invalid; +} + static void configure_gfx1151_dspark_mmvq_default(int gpu) { #if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) if (!env_flag_enabled("DFLASH_DS4_SPEC") || @@ -394,6 +417,42 @@ static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, return true; } +// Safe adaptive cache budget after dense and statically-hot weights have been +// allocated on the same device. This is the Strix-only counterpart of the +// separate expert-GPU budget: reserve the not-yet-created KV cache and runtime +// headroom, then expose only genuine leftover memory to the common SSD tier. +static size_t compute_single_device_stream_cache_budget( + const DeepSeek4Weights & w, + const Ds4ExpertMemoryInfo & placed, + int gpu, + int max_ctx) { + size_t gpu_free = 0; + size_t gpu_total = 0; + ggml_backend_cuda_get_device_memory(gpu, &gpu_free, &gpu_total); + if (gpu_total == 0) return 0; + + const uint64_t kv_bytes = estimate_ds4_cache_bytes(w, max_ctx); + const uint64_t runtime_reserve = + 256ULL * 1024 * 1024 + 512ULL * 1024 * 1024; + uint64_t available = (uint64_t) gpu_free > kv_bytes + runtime_reserve + ? (uint64_t) gpu_free - kv_bytes - runtime_reserve : 0; + + // On a single device, DFLASH_EXPERT_BUDGET_MB is a cap for all routed + // residency, not a second cache allowance on top of the static hot set. + if (const char * cap_env = std::getenv("DFLASH_EXPERT_BUDGET_MB")) { + const uint64_t cap = + (uint64_t) std::max(0, std::atoi(cap_env)) * 1024ULL * 1024ULL; + if (cap > 0) { + const uint64_t cap_remaining = cap > placed.hot_bytes + ? cap - placed.hot_bytes : 0; + available = std::min(available, cap_remaining); + } + } + available = std::min(available, placed.cold_bytes); + return available > (uint64_t) std::numeric_limits::max() + ? std::numeric_limits::max() : (size_t) available; +} + static MoeHybridConfig make_ds4_parent_worker_cfg(const DeepSeek4Weights & w) { MoeHybridConfig cfg; cfg.n_embd = w.n_embd; @@ -476,8 +535,10 @@ bool DeepSeek4Backend::load_model() { // disable the requested split before the TP runtime can initialize. const bool force_full = env_flag_enabled("DFLASH_DS4_FORCE_FULL_LOAD"); const bool heterogeneous_tp = env_flag_enabled("DFLASH_DS4_MOE_TP"); + const bool explicit_ssd_capacity = + moe_nvme_cold_tier_mode() == MoeNvmeColdTierMode::Enabled; const bool need_monolithic = - requires_monolithic_model() && !heterogeneous_tp; + requires_monolithic_model() && !heterogeneous_tp && !explicit_ssd_capacity; if (target_backend == PlacementBackend::Hip && (force_full || need_monolithic)) { std::fprintf(stderr, @@ -791,6 +852,18 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & const bool all_cold = env_flag_enabled("DFLASH_DS4_MOE_TP_ALL_COLD"); int hot_per_layer = all_cold ? 0 : budget.max_hot_per_layer; + if (!all_cold && + moe_nvme_cold_tier_mode() == MoeNvmeColdTierMode::Enabled && + hot_per_layer >= w.n_expert) { + // `on` is also a qualification switch. Keep one exact expert per + // layer in the SSD tier when the whole model would otherwise fit, so + // a single-Strix user can exercise the real path without inventing a + // fragile machine-specific memory cap. `auto` still prefers full + // residency whenever it fits. + hot_per_layer = std::max(0, w.n_expert - 1); + std::fprintf(stderr, + "[deepseek4] forced SSD tier: reserving one cold expert per layer\n"); + } if (all_cold) { std::fprintf(stderr, "[deepseek4-moe-tp] all routed experts assigned to the cold backend\n"); @@ -859,8 +932,23 @@ bool DeepSeek4Backend::init_hybrid_model() { auto hybrid = std::make_shared(); MoeHybridConfig hybrid_cfg = make_ds4_parent_worker_cfg(w_); size_t nvme_device_cache_bytes = 0; + const bool tp_requested = env_flag_enabled("DFLASH_DS4_MOE_TP"); const bool inprocess_tp = - env_flag_enabled("DFLASH_DS4_MOE_TP") && ds4_inprocess_moe_tp_enabled(); + tp_requested && ds4_inprocess_moe_tp_enabled(); + const bool external_tp = tp_requested && !inprocess_tp; + MoeNvmeColdTierMode nvme_mode = moe_nvme_cold_tier_mode(); + if (nvme_mode == MoeNvmeColdTierMode::Invalid) { + std::fprintf(stderr, + "[deepseek4] ignoring invalid DFLASH_MOE_NVME_COLD_TIER; using auto\n"); + nvme_mode = MoeNvmeColdTierMode::Auto; + } + + // A partially resident single-GPU model streams its non-resident experts + // on the same device. In-process TP instead streams them on the expert GPU; + // external TP leaves them to the remote worker and does not start a local + // SSD service. `off` explicitly selects the older materialized CPU tail. + bool stream_cold = !external_tp && + nvme_mode != MoeNvmeColdTierMode::Disabled; if (inprocess_tp) { const int expert_gpu = ds4_moe_tp_gpu(cfg_.device.gpu); if (expert_gpu == cfg_.device.gpu) { @@ -883,18 +971,10 @@ bool DeepSeek4Backend::init_hybrid_model() { } hybrid_cfg.cold_expert_backend = MoeHybridColdBackend::Gpu; - // Third tier: R9700 static-hot -> Strix adaptive-warm -> SSD misses. - // Keep the established fully-resident path whenever it fits. In auto - // mode, switch only when the cold stack plus a safety reserve exceeds - // currently available Strix memory. An explicit override is useful - // for qualification on models that already fit. - bool stream_cold = false; - const char * nvme_mode = std::getenv("DFLASH_MOE_NVME_COLD_TIER"); - if (nvme_mode && (std::strcmp(nvme_mode, "1") == 0 || - std::strcmp(nvme_mode, "on") == 0 || - std::strcmp(nvme_mode, "true") == 0)) { - stream_cold = true; - } else if (!nvme_mode || std::strcmp(nvme_mode, "auto") == 0) { + // On a separate expert GPU, auto mode retains the established fully + // resident path whenever the cold stack fits after a conservative + // reserve. Explicit on/off remain authoritative. + if (nvme_mode == MoeNvmeColdTierMode::Auto) { Ds4ExpertMemoryInfo info; std::string memory_error; size_t expert_free = 0; @@ -914,15 +994,9 @@ bool DeepSeek4Backend::init_hybrid_model() { gib(nvme_device_cache_bytes), stream_cold ? "ssd-stream" : "resident"); } - } else if (std::strcmp(nvme_mode, "0") != 0 && - std::strcmp(nvme_mode, "off") != 0 && - std::strcmp(nvme_mode, "false") != 0) { - std::fprintf(stderr, - "[deepseek4] ignoring invalid DFLASH_MOE_NVME_COLD_TIER=%s\n", - nvme_mode); } - hybrid_cfg.materialize_cold_experts = !stream_cold; } + hybrid_cfg.materialize_cold_experts = !stream_cold && !external_tp; if (!build_deepseek4_moe_hybrid_storage_from_file_with_mmap( cfg_.model_path, backend_, w_, moe_placement_, &hybrid_cfg, *hybrid, &err, expert_backend_)) { @@ -934,12 +1008,39 @@ bool DeepSeek4Backend::init_hybrid_model() { return false; } - if (hybrid->has_mmap() && !hybrid->materialized_cold_experts) { + if (stream_cold && !inprocess_tp && + !std::getenv("DFLASH_MOE_NVME_DEVICE_CACHE_MB")) { + Ds4ExpertMemoryInfo placed; + std::string memory_error; + if (compute_ds4_expert_memory_info(w_, &moe_placement_, + placed, &memory_error)) { + nvme_device_cache_bytes = compute_single_device_stream_cache_budget( + w_, placed, cfg_.device.gpu, max_ctx); + size_t device_free = 0; + size_t device_total = 0; + ggml_backend_cuda_get_device_memory( + cfg_.device.gpu, &device_free, &device_total); + std::fprintf(stderr, + "[deepseek4] single-device SSD tier: device=%d free=%.2f GiB " + "cold=%.2f GiB adaptive-cache=%.2f GiB\n", + cfg_.device.gpu, gib(device_free), gib(placed.cold_bytes), + gib(nvme_device_cache_bytes)); + } + } + + if (!external_tp && hybrid->has_mmap() && + !hybrid->materialized_cold_experts) { size_t max_expert_bytes = 0; - for (const auto & layer : hybrid->layers) { - const size_t per_expert_bytes = layer.fused_gate_up + for (size_t il = 0; il < hybrid->layers.size(); ++il) { + const auto & layer = hybrid->layers[il]; + size_t per_expert_bytes = layer.fused_gate_up ? layer.gate_up_expert_bytes + layer.down_expert_bytes : layer.gate_expert_bytes + layer.up_expert_bytes + layer.down_expert_bytes; + if (il < hybrid->layer_regions.size() && + hybrid->layer_regions[il].expert_major.enabled) { + per_expert_bytes = + hybrid->layer_regions[il].expert_major.expert_stride; + } max_expert_bytes = std::max(max_expert_bytes, per_expert_bytes); } if (max_expert_bytes == 0) { @@ -960,9 +1061,10 @@ bool DeepSeek4Backend::init_hybrid_model() { } std::fprintf(stderr, "[deepseek4] cold-expert SSD engine ready: io=%s pinned=%.1f MiB " - "strix_cache=%.1f MiB slots=%d\n", + "device=%d device_cache=%.1f MiB slots=%d\n", stream_engine_.io_backend_name(), stream_engine_.pinned_bytes() / 1024.0 / 1024.0, + inprocess_tp ? ds4_moe_tp_gpu(cfg_.device.gpu) : cfg_.device.gpu, stream_engine_.device_cache_bytes() / 1024.0 / 1024.0, stream_engine_.device_slot_count()); } @@ -970,8 +1072,10 @@ bool DeepSeek4Backend::init_hybrid_model() { moe_hybrid_ = std::move(hybrid); w_.moe_hybrid = true; const int total_cold = w_.n_layer * w_.n_expert - moe_placement_.total_hot; - const char * cold_backend = - moe_hybrid_->cold_backend_kind == MoeHybridColdBackend::Gpu ? "gpu" : "cpu"; + const char * cold_backend = stream_engine_.is_bound() + ? "ssd" : (external_tp ? "remote" : + (moe_hybrid_->cold_backend_kind == MoeHybridColdBackend::Gpu + ? "gpu" : "cpu")); std::fprintf(stderr, "[deepseek4] hybrid experts ready: hot=%d cold=%d cold_backend=%s%s\n", moe_placement_.total_hot, total_cold, cold_backend, ""); return true; diff --git a/server/test/bench_kimi_k3_hetero.cpp b/server/test/bench_kimi_k3_hetero.cpp index bfc98bd5d..581498a41 100644 --- a/server/test/bench_kimi_k3_hetero.cpp +++ b/server/test/bench_kimi_k3_hetero.cpp @@ -61,115 +61,6 @@ double gib(uint64_t bytes) { return (double) bytes / (1024.0 * 1024.0 * 1024.0); } -class PersistentKimiExpertGraph { -public: - ~PersistentKimiExpertGraph() { destroy(); } - - bool init(ggml_backend_t backend, std::string & error) { - backend_ = backend; - ggml_init_params params{}; - params.mem_size = 16 * 1024 * 1024; - params.no_alloc = true; - ctx_ = ggml_init(params); - if (!ctx_) { - error = "ggml_init failed for Kimi expert graph"; - return false; - } - - gate_ = ggml_new_tensor_2d( - ctx_, GGML_TYPE_IQ1_S, kKimiLatent, kKimiExpertFf); - up_ = ggml_new_tensor_2d( - ctx_, GGML_TYPE_IQ1_S, kKimiLatent, kKimiExpertFf); - down_ = ggml_new_tensor_2d( - ctx_, GGML_TYPE_IQ1_S, kKimiExpertFf, kKimiLatent); - input_ = ggml_new_tensor_2d(ctx_, GGML_TYPE_F32, kKimiLatent, 1); - ggml_set_input(gate_); - ggml_set_input(up_); - ggml_set_input(down_); - ggml_set_input(input_); - - ggml_tensor * gate_value = ggml_mul_mat(ctx_, gate_, input_); - ggml_tensor * up_value = ggml_mul_mat(ctx_, up_, input_); - - // SiTU(g, u) = beta*tanh(g/beta)*sigmoid(g) - // * linear_beta*tanh(u/linear_beta). - ggml_tensor * activated = ggml_scale(ctx_, gate_value, 1.0f / kSituBeta); - activated = ggml_tanh(ctx_, activated); - activated = ggml_scale(ctx_, activated, kSituBeta); - activated = ggml_mul(ctx_, activated, ggml_sigmoid(ctx_, gate_value)); - up_value = ggml_scale(ctx_, up_value, 1.0f / kSituLinearBeta); - up_value = ggml_tanh(ctx_, up_value); - up_value = ggml_scale(ctx_, up_value, kSituLinearBeta); - activated = ggml_mul(ctx_, activated, up_value); - output_ = ggml_mul_mat(ctx_, down_, activated); - ggml_set_output(output_); - - graph_ = ggml_new_graph_custom(ctx_, 256, false); - ggml_build_forward_expand(graph_, output_); - buffer_ = ggml_backend_alloc_ctx_tensors(ctx_, backend_); - if (!buffer_) { - error = "device allocation failed for Kimi expert graph"; - return false; - } - - std::vector input((size_t) kKimiLatent); - for (size_t i = 0; i < input.size(); ++i) { - input[i] = 0.01f * std::sin((float) i * 0.013f); - } - ggml_backend_tensor_set( - input_, input.data(), 0, input.size() * sizeof(float)); - return true; - } - - bool launch(const MoeHybridStreamEngine & engine, std::string & error) { - if (!ctx_ || !graph_ || !backend_) { - error = "Kimi expert graph is not initialized"; - return false; - } - if (engine.scratch_gate_bytes() != ggml_nbytes(gate_) || - engine.scratch_up_bytes() != ggml_nbytes(up_) || - engine.scratch_down_bytes() != ggml_nbytes(down_)) { - error = "streamed Kimi expert byte layout does not match IQ1_S graph"; - return false; - } - gate_->data = const_cast(engine.scratch_gate_data()); - up_->data = const_cast(engine.scratch_up_data()); - down_->data = const_cast(engine.scratch_down_data()); - if (ggml_backend_graph_compute_async(backend_, graph_) != - GGML_STATUS_SUCCESS) { - error = "Kimi expert graph launch failed"; - return false; - } - return true; - } - - void destroy() { - if (backend_) ggml_backend_synchronize(backend_); - if (buffer_) { - ggml_backend_buffer_free(buffer_); - buffer_ = nullptr; - } - if (ctx_) { - ggml_free(ctx_); - ctx_ = nullptr; - } - graph_ = nullptr; - backend_ = nullptr; - gate_ = up_ = down_ = input_ = output_ = nullptr; - } - -private: - ggml_backend_t backend_ = nullptr; - ggml_context * ctx_ = nullptr; - ggml_cgraph * graph_ = nullptr; - ggml_backend_buffer_t buffer_ = nullptr; - ggml_tensor * gate_ = nullptr; - ggml_tensor * up_ = nullptr; - ggml_tensor * down_ = nullptr; - ggml_tensor * input_ = nullptr; - ggml_tensor * output_ = nullptr; -}; - std::vector route_for(int token, int layer, int top_k, bool repeat) { std::vector population((size_t) kKimiExperts); std::iota(population.begin(), population.end(), 0); @@ -289,13 +180,24 @@ int main(int argc, char ** argv) { return 1; } - PersistentKimiExpertGraph expert_graph; - if (compute_arg && !expert_graph.init(backend, error)) { - std::fprintf(stderr, "Kimi graph initialization failed: %s\n", error.c_str()); - engine.destroy(); - ggml_backend_free(backend); - return 1; + MoeStreamExpertSpec expert_spec; + expert_spec.input_dim = (int) kKimiLatent; + expert_spec.intermediate_dim = (int) kKimiExpertFf; + expert_spec.output_dim = (int) kKimiLatent; + expert_spec.gate_type = GGML_TYPE_IQ1_S; + expert_spec.up_type = GGML_TYPE_IQ1_S; + expert_spec.down_type = GGML_TYPE_IQ1_S; + expert_spec.gated_activation = MoeGatedActivation::Situ; + expert_spec.situ_beta = kSituBeta; + expert_spec.situ_linear_beta = kSituLinearBeta; + + std::vector model_input((size_t) kKimiLatent); + for (size_t i = 0; i < model_input.size(); ++i) { + model_input[i] = 0.01f * std::sin((float) i * 0.013f); } + std::vector route_weights( + (size_t) top_k_arg, 1.0f / (float) top_k_arg); + std::vector routed_output; const uint64_t accesses = tokens_arg * layers_arg * top_k_arg; @@ -304,6 +206,25 @@ int main(int argc, char ** argv) { for (int layer = 0; layer < (int) layers_arg; ++layer) { const std::vector experts = route_for( token, layer, (int) top_k_arg, repeat_arg != 0); + if (compute_arg) { + MoeStreamRouteBatch batch; + batch.layer = layer; + batch.n_expert = kKimiExperts; + batch.top_k = (int) top_k_arg; + batch.n_tokens = 1; + batch.inputs = model_input.data(); + batch.selected_ids = experts.data(); + batch.selected_weights = route_weights.data(); + if (!eval_moe_streamed_experts( + engine, expert_spec, batch, routed_output, &error)) { + std::fprintf(stderr, + "common streamed-expert evaluation failed: %s\n", + error.c_str()); + return 1; + } + continue; + } + engine.request_experts( layer, experts.data(), (int) experts.size(), MoeNvmePriority::Demand); @@ -319,12 +240,7 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "expert activation failed: %s\n", error.c_str()); return 1; } - if (compute_arg && !expert_graph.launch(engine, error)) { - std::fprintf(stderr, "expert compute failed: %s\n", error.c_str()); - return 1; - } - - // Disk/H2D for expert N+1 overlaps expert N's persistent graph. + // Disk/H2D for expert N+1 overlaps bookkeeping for expert N. if (i + 1 < experts.size()) { int next_slot = -1; if (!engine.stage_expert_cached_async( @@ -335,7 +251,6 @@ int main(int argc, char ** argv) { } staged_slot = next_slot; } - if (compute_arg) ggml_backend_synchronize(backend); engine.release_device_slot(current_slot); } } @@ -345,6 +260,7 @@ int main(int argc, char ** argv) { const double seconds = std::chrono::duration(end - begin).count(); const MoeNvmeStats stats = engine.io_stats(); + const MoeStreamComputeStats compute_stats = engine.compute_stats(); const double misses = expert_bytes > 0 ? (double) stats.payload_bytes / (double) expert_bytes : 0.0; const double hit_rate = accesses > 0 @@ -369,12 +285,15 @@ int main(int argc, char ** argv) { seconds > 0 ? (double) accesses / seconds : 0.0); std::printf( "ssd_payload_gib=%.6f physical_gib=%.6f pipeline_gib_s=%.6f " - "estimated_device_cache_hit=%.4f cache_gib=%.3f io_errors=%" PRIu64 "\n", + "estimated_device_cache_hit=%.4f cache_gib=%.3f io_errors=%" PRIu64 + " graph_builds=%" PRIu64 " graph_hits=%" PRIu64 + " graph_launches=%" PRIu64 "\n", gib(stats.payload_bytes), gib(stats.physical_bytes), seconds > 0 ? gib(stats.payload_bytes) / seconds : 0.0, - hit_rate, gib(engine.device_cache_bytes()), stats.errors); + hit_rate, gib(engine.device_cache_bytes()), stats.errors, + compute_stats.graph_builds, compute_stats.graph_cache_hits, + compute_stats.graph_launches); - expert_graph.destroy(); engine.destroy(); ggml_backend_free(backend); #if defined(_WIN32) diff --git a/server/test/test_moe_nvme_scheduler.cpp b/server/test/test_moe_nvme_scheduler.cpp index e85014a81..75f72afd0 100644 --- a/server/test/test_moe_nvme_scheduler.cpp +++ b/server/test/test_moe_nvme_scheduler.cpp @@ -96,6 +96,7 @@ void verify_lease(const MoeNvmeLease & lease, int layer, int expert) { NVME_REQUIRE(layout.key.expert == expert); const int expected_spans = layer == 0 ? 3 : 2; NVME_REQUIRE(layout.span_count == expected_spans); + NVME_REQUIRE(layout.component_count == expected_spans); for (int tensor = 0; tensor < layout.span_count; ++tensor) { const MoeExpertIoSpan & span = layout.spans[tensor]; const uint8_t * payload = lease.data() + span.buffer_offset; @@ -124,6 +125,7 @@ TEST_CASE(MoeNvmeSchedulerFixture, exact_layout_bounds_and_alignment) { NVME_REQUIRE(make_moe_expert_io_layout( 0, 3, model.regions[0], model.file.size(), 4096, layout, &err)); NVME_REQUIRE(layout.span_count == 3); + NVME_REQUIRE(layout.component_count == 3); NVME_REQUIRE(layout.payload_bytes == model.regions[0].expert_bytes_gate + model.regions[0].expert_bytes_up + @@ -135,6 +137,75 @@ TEST_CASE(MoeNvmeSchedulerFixture, exact_layout_bounds_and_alignment) { NVME_REQUIRE(!err.empty()); } +TEST_CASE(MoeNvmeSchedulerFixture, expert_major_layout_uses_one_read_without_changing_components) { + constexpr int experts = SyntheticModel::kExperts; + constexpr size_t gate_bytes = 4093; + constexpr size_t up_bytes = 6141; + constexpr size_t down_bytes = 8189; + constexpr size_t gate_offset = 0; + constexpr size_t up_offset = 4352; + constexpr size_t down_offset = 10752; + constexpr size_t stride = 19456; + constexpr size_t file_offset = 4096; + + std::vector file(file_offset + stride * experts + 4096, 0xa5); + LayerExpertRegions layer; + layer.expert_bytes_gate = gate_bytes; + layer.expert_bytes_up = up_bytes; + layer.expert_bytes_down = down_bytes; + layer.expert_major.enabled = true; + layer.expert_major.expert_stride = stride; + layer.expert_major.gate_offset = gate_offset; + layer.expert_major.up_offset = up_offset; + layer.expert_major.down_offset = down_offset; + layer.expert_major.experts = {file_offset, stride * experts, 0}; + for (int expert = 0; expert < experts; ++expert) { + const size_t base = file_offset + (size_t) expert * stride; + for (size_t i = 0; i < gate_bytes; ++i) { + file[base + gate_offset + i] = expected_byte(0, 0, expert, i); + } + for (size_t i = 0; i < up_bytes; ++i) { + file[base + up_offset + i] = expected_byte(0, 1, expert, i); + } + for (size_t i = 0; i < down_bytes; ++i) { + file[base + down_offset + i] = expected_byte(0, 2, expert, i); + } + } + + MoeExpertIoLayout layout; + std::string err; + NVME_REQUIRE(make_moe_expert_io_layout( + 0, 5, layer, file.size(), 4096, layout, &err)); + NVME_REQUIRE(layout.span_count == 1); + NVME_REQUIRE(layout.component_count == 3); + NVME_REQUIRE(layout.payload_bytes == stride); + NVME_REQUIRE(layout.component(MoeExpertComponentKind::Gate)->device_offset == gate_offset); + NVME_REQUIRE(layout.component(MoeExpertComponentKind::Up)->device_offset == up_offset); + NVME_REQUIRE(layout.component(MoeExpertComponentKind::Down)->device_offset == down_offset); + + MoeNvmeConfig config; + config.backend = MoeNvmeBackend::Mmap; + config.direct_io = MoeNvmeDirectMode::Disabled; + config.host_slots = 4; + MoeNvmeScheduler scheduler; + NVME_REQUIRE(scheduler.init(config, stride, aligned_allocate, + aligned_free, nullptr, &err)); + NVME_REQUIRE(scheduler.bind_source( + {file.data(), file.size(), -1}, {layer}, &err)); + MoeNvmeLease lease; + NVME_REQUIRE(scheduler.acquire(0, 5, lease, &err)); + const MoeExpertIoSpan & record = lease.layout().spans[0]; + const uint8_t * payload = lease.data() + record.buffer_offset; + NVME_REQUIRE(payload[gate_offset] == expected_byte(0, 0, 5, 0)); + NVME_REQUIRE(payload[up_offset + up_bytes / 2] == + expected_byte(0, 1, 5, up_bytes / 2)); + NVME_REQUIRE(payload[down_offset + down_bytes - 1] == + expected_byte(0, 2, 5, down_bytes - 1)); + lease.reset(); + NVME_REQUIRE(scheduler.stats().read_ops == 1); + NVME_REQUIRE(scheduler.stats().errors == 0); +} + TEST_CASE(MoeNvmeSchedulerFixture, async_exact_reads_dedupe_and_cache) { SyntheticModel model; MoeNvmeConfig config; diff --git a/server/test/test_moe_stream_compute.cpp b/server/test/test_moe_stream_compute.cpp new file mode 100644 index 000000000..b8006e5e0 --- /dev/null +++ b/server/test/test_moe_stream_compute.cpp @@ -0,0 +1,311 @@ +#include "CppUnitTestFramework.hpp" +#include "common/moe_hybrid_stream.h" + +#include "ggml-cuda.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +using namespace dflash::common; + +#define STREAM_REQUIRE(cond) do { \ + if (!(cond)) throw std::runtime_error(std::string(__FILE__) + ":" + \ + std::to_string(__LINE__) + ": " + #cond); \ +} while (0) + +namespace { + +struct MoeStreamComputeFixture {}; + +constexpr int kExperts = 3; +constexpr int kInput = 16; +constexpr int kFf = 24; +constexpr int kOutput = 12; +constexpr int kTokens = 2; +constexpr int kTopK = 2; + +struct TempFile { + int fd = -1; + + explicit TempFile(const std::vector & bytes) { + char path[] = "/tmp/luce-moe-stream-XXXXXX"; + fd = ::mkstemp(path); + if (fd < 0) throw std::runtime_error("mkstemp failed"); + (void) ::unlink(path); + size_t done = 0; + while (done < bytes.size()) { + const ssize_t wrote = ::pwrite( + fd, bytes.data() + done, bytes.size() - done, (off_t) done); + if (wrote <= 0) throw std::runtime_error("pwrite failed"); + done += (size_t) wrote; + } + } + + ~TempFile() { + if (fd >= 0) ::close(fd); + } +}; + +float gate_value(int expert, int row, int column) { + return 0.08f * std::sin( + 0.17f * (float) (1 + expert * 11 + row * 5 + column)); +} + +float up_value(int expert, int row, int column) { + return 0.07f * std::cos( + 0.13f * (float) (3 + expert * 7 + row * 3 + column)); +} + +float down_value(int expert, int row, int column) { + return 0.06f * std::sin( + 0.11f * (float) (5 + expert * 13 + row * 2 + column)); +} + +void fill_weights(std::vector & gate, + std::vector & up, + std::vector & down) { + gate.resize((size_t) kExperts * kFf * kInput); + up.resize(gate.size()); + down.resize((size_t) kExperts * kOutput * kFf); + for (int expert = 0; expert < kExperts; ++expert) { + for (int row = 0; row < kFf; ++row) { + for (int column = 0; column < kInput; ++column) { + const size_t i = ((size_t) expert * kFf + row) * kInput + column; + gate[i] = gate_value(expert, row, column); + up[i] = up_value(expert, row, column); + } + } + for (int row = 0; row < kOutput; ++row) { + for (int column = 0; column < kFf; ++column) { + const size_t i = ((size_t) expert * kOutput + row) * kFf + column; + down[i] = down_value(expert, row, column); + } + } + } +} + +void append_at(std::vector & file, size_t offset, + const float * values, size_t count) { + STREAM_REQUIRE(offset <= file.size()); + STREAM_REQUIRE(count * sizeof(float) <= file.size() - offset); + std::memcpy(file.data() + offset, values, count * sizeof(float)); +} + +struct ModelBytes { + std::vector file; + LayerExpertRegions regions; + size_t slot_bytes = 0; +}; + +ModelBytes make_model_bytes(bool expert_major, + const std::vector & gate, + const std::vector & up, + const std::vector & down) { + const size_t gate_bytes = (size_t) kInput * kFf * sizeof(float); + const size_t up_bytes = gate_bytes; + const size_t down_bytes = (size_t) kFf * kOutput * sizeof(float); + ModelBytes model; + model.regions.expert_bytes_gate = gate_bytes; + model.regions.expert_bytes_up = up_bytes; + model.regions.expert_bytes_down = down_bytes; + + if (!expert_major) { + const size_t gate_stack = gate_bytes * kExperts; + const size_t up_stack = up_bytes * kExperts; + const size_t down_stack = down_bytes * kExperts; + model.file.resize(gate_stack + up_stack + down_stack); + model.regions.gate_exps = {0, gate_stack}; + model.regions.up_exps = {gate_stack, up_stack}; + model.regions.down_exps = {gate_stack + up_stack, down_stack}; + std::memcpy(model.file.data(), gate.data(), gate_stack); + std::memcpy(model.file.data() + gate_stack, up.data(), up_stack); + std::memcpy(model.file.data() + gate_stack + up_stack, + down.data(), down_stack); + model.slot_bytes = gate_bytes + up_bytes + down_bytes; + return model; + } + + constexpr size_t kGap = 256; + const size_t gate_offset = 0; + const size_t up_offset = gate_bytes + kGap; + const size_t down_offset = up_offset + up_bytes + kGap; + const size_t stride = down_offset + down_bytes; + model.file.assign(stride * kExperts, 0); + model.regions.expert_major.enabled = true; + model.regions.expert_major.experts = {0, model.file.size()}; + model.regions.expert_major.expert_stride = stride; + model.regions.expert_major.gate_offset = gate_offset; + model.regions.expert_major.up_offset = up_offset; + model.regions.expert_major.down_offset = down_offset; + for (int expert = 0; expert < kExperts; ++expert) { + const size_t base = (size_t) expert * stride; + append_at(model.file, base + gate_offset, + gate.data() + (size_t) expert * kFf * kInput, + (size_t) kFf * kInput); + append_at(model.file, base + up_offset, + up.data() + (size_t) expert * kFf * kInput, + (size_t) kFf * kInput); + append_at(model.file, base + down_offset, + down.data() + (size_t) expert * kOutput * kFf, + (size_t) kOutput * kFf); + } + model.slot_bytes = stride; + return model; +} + +std::vector cpu_reference( + const std::vector & gate, + const std::vector & up, + const std::vector & down, + const std::vector & input, + const int32_t * ids, + const float * weights) { + constexpr float gate_scale = 0.8f; + constexpr float up_scale = 1.1f; + constexpr float down_scale = 0.9f; + constexpr float beta = 4.0f; + constexpr float linear_beta = 25.0f; + std::vector output((size_t) kTokens * kOutput, 0.0f); + std::vector activated(kFf); + for (int token = 0; token < kTokens; ++token) { + for (int rank = 0; rank < kTopK; ++rank) { + const int expert = ids[token * kTopK + rank]; + for (int row = 0; row < kFf; ++row) { + float g = 0.0f; + float u = 0.0f; + for (int column = 0; column < kInput; ++column) { + const size_t wi = + ((size_t) expert * kFf + row) * kInput + column; + const float x = input[(size_t) token * kInput + column]; + g += gate[wi] * x; + u += up[wi] * x; + } + g *= gate_scale; + u *= up_scale; + const float nonlinear = + beta * std::tanh(g / beta) / (1.0f + std::exp(-g)); + const float linear = linear_beta * std::tanh(u / linear_beta); + activated[(size_t) row] = nonlinear * linear; + } + for (int row = 0; row < kOutput; ++row) { + float value = 0.0f; + for (int column = 0; column < kFf; ++column) { + const size_t wi = + ((size_t) expert * kOutput + row) * kFf + column; + value += down[wi] * activated[(size_t) column]; + } + output[(size_t) token * kOutput + row] += + weights[token * kTopK + rank] * down_scale * value; + } + } + } + return output; +} + +void run_layout_case(ggml_backend_t backend, bool expert_major) { + std::vector gate; + std::vector up; + std::vector down; + fill_weights(gate, up, down); + ModelBytes model = make_model_bytes(expert_major, gate, up, down); + TempFile file(model.file); + + MoeHybridStorage storage; + storage.mmap_size = model.file.size(); + storage.mmap_fd = ::dup(file.fd); + STREAM_REQUIRE(storage.mmap_fd >= 0); + storage.layer_regions.push_back(model.regions); + + MoeStreamConfig config; + config.device_slots = 2; + config.device_cache_bytes = 0; + config.graph_cache_entries = 4; + config.nvme.backend = MoeNvmeBackend::ThreadPool; + config.nvme.direct_io = MoeNvmeDirectMode::Disabled; + config.nvme.host_slots = 6; + config.nvme.io_threads = 2; + + MoeHybridStreamEngine engine; + std::string error; + STREAM_REQUIRE(engine.init( + backend, model.slot_bytes, storage, config, &error)); + + MoeStreamExpertSpec spec; + spec.input_dim = kInput; + spec.intermediate_dim = kFf; + spec.output_dim = kOutput; + spec.gate_type = GGML_TYPE_F32; + spec.up_type = GGML_TYPE_F32; + spec.down_type = GGML_TYPE_F32; + spec.gated_activation = MoeGatedActivation::Situ; + spec.gate_scale = 0.8f; + spec.up_scale = 1.1f; + spec.down_scale = 0.9f; + + std::vector input((size_t) kTokens * kInput); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = 0.12f * std::sin(0.07f * (float) (i + 1)); + } + const int32_t ids[kTokens * kTopK] = {2, 0, 1, 2}; + const float weights[kTokens * kTopK] = {0.65f, 0.35f, 0.55f, 0.45f}; + MoeStreamRouteBatch batch; + batch.layer = 0; + batch.n_expert = kExperts; + batch.top_k = kTopK; + batch.n_tokens = kTokens; + batch.inputs = input.data(); + batch.selected_ids = ids; + batch.selected_weights = weights; + + const std::vector expected = + cpu_reference(gate, up, down, input, ids, weights); + std::vector actual; + STREAM_REQUIRE(eval_moe_streamed_experts( + engine, spec, batch, actual, &error)); + STREAM_REQUIRE(actual.size() == expected.size()); + for (size_t i = 0; i < actual.size(); ++i) { + const float tolerance = 2.0e-5f + 2.0e-4f * std::fabs(expected[i]); + STREAM_REQUIRE(std::fabs(actual[i] - expected[i]) <= tolerance); + } + + const MoeStreamComputeStats first = engine.compute_stats(); + STREAM_REQUIRE(first.graph_builds == 2); + STREAM_REQUIRE(first.graph_launches == 3); + STREAM_REQUIRE(eval_moe_streamed_experts( + engine, spec, batch, actual, &error)); + const MoeStreamComputeStats second = engine.compute_stats(); + STREAM_REQUIRE(second.graph_builds == first.graph_builds); + STREAM_REQUIRE(second.graph_cache_hits > first.graph_cache_hits); + STREAM_REQUIRE(second.graph_launches == 6); + engine.destroy(); +} + +} // namespace + +TEST_CASE(MoeStreamComputeFixture, persistent_graph_matches_cpu_for_both_layouts) { + int device = 0; + if (const char * value = std::getenv("DFLASH_TEST_GPU")) { + device = std::max(0, std::atoi(value)); + } + if (device >= ggml_backend_cuda_get_device_count()) { + std::fprintf(stderr, "skip: requested CUDA/HIP device is unavailable\n"); + return; + } + ggml_backend_t backend = ggml_backend_cuda_init(device); + if (!backend) { + std::fprintf(stderr, "skip: no CUDA/HIP backend available\n"); + return; + } + run_layout_case(backend, false); + run_layout_case(backend, true); + ggml_backend_free(backend); +} From 49b2780c7c5a01029ccf291dca0efa664680fb35 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:19:00 +0200 Subject: [PATCH 03/20] perf(moe): reuse streamed expert host buffers --- server/src/common/moe_hybrid_stream.cpp | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/server/src/common/moe_hybrid_stream.cpp b/server/src/common/moe_hybrid_stream.cpp index 0bc86a9c1..7f9c42697 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -1614,6 +1614,12 @@ bool eval_moe_streamed_experts( return true; }; + struct TokenHit { int token; float weight; }; + std::vector hits; + hits.reserve((size_t) batch.n_tokens); + std::vector compact_input; + std::vector result; + for (size_t expert_index = 0; expert_index < unique_experts.size(); ++expert_index) { const int current_slot = staged_slot; @@ -1622,9 +1628,7 @@ bool eval_moe_streamed_experts( engine.release_device_slot(current_slot); }; - struct TokenHit { int token; float weight; }; - std::vector hits; - hits.reserve((size_t) batch.n_tokens); + hits.clear(); const int32_t expert = unique_experts[expert_index]; for (int token = 0; token < batch.n_tokens; ++token) { float combined_weight = 0.0f; @@ -1655,7 +1659,7 @@ bool eval_moe_streamed_experts( release_current(); return false; } - std::vector compact_input(input_values); + compact_input.resize(input_values); for (size_t i = 0; i < hits.size(); ++i) { const float * src = batch.inputs + (size_t) hits[i].token * (size_t) spec.input_dim; @@ -1694,7 +1698,6 @@ bool eval_moe_streamed_experts( staged_slot = next_slot; } - std::vector result; if (!graph->finish(result, err)) { release_current(); return false; From ad29ab104a66335e26cae6614d51e5206cc7f8c8 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:20:32 +0200 Subject: [PATCH 04/20] feat(kimi): add native K3 target backend --- server/CMakeLists.txt | 16 + server/src/common/backend_factory.cpp | 15 + server/src/common/model_capabilities.h | 1 + server/src/kimi_k3/kimi_k3_backend.cpp | 221 +++++++++++ server/src/kimi_k3/kimi_k3_backend.h | 59 +++ server/src/kimi_k3/kimi_k3_graph.cpp | 496 +++++++++++++++++++++++++ server/src/kimi_k3/kimi_k3_internal.h | 159 ++++++++ server/src/kimi_k3/kimi_k3_loader.cpp | 326 ++++++++++++++++ server/test/smoke_kimi_k3_forward.cpp | 70 ++++ server/test/test_feature_gate.cpp | 11 +- 10 files changed, 1370 insertions(+), 4 deletions(-) create mode 100644 server/src/kimi_k3/kimi_k3_backend.cpp create mode 100644 server/src/kimi_k3/kimi_k3_backend.h create mode 100644 server/src/kimi_k3/kimi_k3_graph.cpp create mode 100644 server/src/kimi_k3/kimi_k3_internal.h create mode 100644 server/src/kimi_k3/kimi_k3_loader.cpp create mode 100644 server/test/smoke_kimi_k3_forward.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index fd1c6ffa2..34237f1cb 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -237,6 +237,7 @@ set(DFLASH27B_SRC_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/src/qwen3 ${CMAKE_CURRENT_SOURCE_DIR}/src/gemma4 ${CMAKE_CURRENT_SOURCE_DIR}/src/deepseek4 + ${CMAKE_CURRENT_SOURCE_DIR}/src/kimi_k3 ${CMAKE_CURRENT_SOURCE_DIR}/src/server ) @@ -269,6 +270,10 @@ add_library(dflash_common STATIC src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp src/deepseek4/deepseek4_dspark.cpp src/deepseek4/deepseek4_dspark_spec.cpp + # Kimi-K3 hybrid KDA/MLA + latent-MoE target arch + src/kimi_k3/kimi_k3_loader.cpp + src/kimi_k3/kimi_k3_graph.cpp + src/kimi_k3/kimi_k3_backend.cpp src/flashprefill_q8.cpp src/kv_cache.cpp src/kv_quant.cpp @@ -988,6 +993,17 @@ if(DFLASH27B_TESTS) target_include_directories(smoke_qwen3_forward PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) target_link_libraries(smoke_qwen3_forward PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/smoke_kimi_k3_forward.cpp") + add_executable(smoke_kimi_k3_forward test/smoke_kimi_k3_forward.cpp) + target_include_directories(smoke_kimi_k3_forward PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(smoke_kimi_k3_forward PRIVATE + dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(smoke_kimi_k3_forward PRIVATE CUDA::cudart) + else() + target_link_libraries(smoke_kimi_k3_forward PRIVATE hip::host) + endif() + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_vs_oracle.cpp") add_executable(test_vs_oracle test/test_vs_oracle.cpp) target_include_directories(test_vs_oracle PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index ab34f2a19..15583e06f 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -13,6 +13,7 @@ #include "gemma4_layer_split_adapter.h" #include "deepseek4_backend.h" #include "deepseek4_layer_split_adapter.h" +#include "kimi_k3_backend.h" #include "layer_split_backend.h" #include "qwen35_layer_split_adapter.h" @@ -94,6 +95,7 @@ DFLASH_CHECK_ARCH("laguna", LagunaBackendArgs, LagunaLayerSplitAdapterCon DFLASH_CHECK_ARCH("qwen3", Qwen3BackendConfig, NoLayerSplitConfig); DFLASH_CHECK_ARCH("gemma4", Gemma4BackendConfig, Gemma4LayerSplitAdapterConfig); DFLASH_CHECK_ARCH("deepseek4", DeepSeek4BackendConfig, DeepSeek4LayerSplitAdapterConfig); +DFLASH_CHECK_ARCH("kimi-k3", KimiK3BackendConfig, NoLayerSplitConfig); #undef DFLASH_CHECK_ARCH #undef DFLASH_CHECK_ARCH_OPTION @@ -436,6 +438,19 @@ std::unique_ptr create_backend( } return backend; + } else if (arch == "kimi-k3") { + KimiK3BackendConfig cfg; + cfg.model_path = args.model_path; + cfg.device = args.device; + cfg.stream_fd = args.stream_fd; + + auto backend = std::make_unique(cfg); + if (!backend->init()) { + std::fprintf(stderr, "[backend_factory] KimiK3Backend init failed\n"); + return nullptr; + } + return backend; + } else { std::fprintf(stderr, "[backend_factory] unsupported architecture: %s\n", arch.c_str()); diff --git a/server/src/common/model_capabilities.h b/server/src/common/model_capabilities.h index 2189fe312..bc956e7aa 100644 --- a/server/src/common/model_capabilities.h +++ b/server/src/common/model_capabilities.h @@ -69,6 +69,7 @@ inline constexpr ArchCapabilities kArchCapabilities[] = { {"qwen3", false, false, true, false, kNever, kNever, kNever, kNever, kNever}, {"gemma4", true, false, false, false, kMono, kNever, kNever, kBoth, kNever}, {"deepseek4", true, false, false, false, kNever, kNever, kNever, kNever, kNever}, + {"kimi-k3", false, false, false, false, kNever, kNever, kNever, kNever, kNever}, }; inline constexpr std::size_t kArchCount = diff --git a/server/src/kimi_k3/kimi_k3_backend.cpp b/server/src/kimi_k3/kimi_k3_backend.cpp new file mode 100644 index 000000000..bf9873273 --- /dev/null +++ b/server/src/kimi_k3/kimi_k3_backend.cpp @@ -0,0 +1,221 @@ +#include "kimi_k3_backend.h" + +#include "common/sampler.h" +#include "dflash27b.h" + +#include "ggml-cuda.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +KimiK3Backend::KimiK3Backend(const KimiK3BackendConfig & cfg) : cfg_(cfg) {} + +KimiK3Backend::~KimiK3Backend() { + shutdown(); +} + +bool KimiK3Backend::init() { + if (!cfg_.model_path) { + std::fprintf(stderr, "[kimi-k3] model path is null\n"); + return false; + } + backend_ = ggml_backend_cuda_init(cfg_.device.primary_gpu()); + if (!backend_) { + std::fprintf(stderr, "[kimi-k3] GPU backend init failed for device %d\n", + cfg_.device.primary_gpu()); + return false; + } + if (!load_kimi_k3_gguf(cfg_.model_path, backend_, weights_)) { + std::fprintf(stderr, "[kimi-k3] model load failed: %s\n", + dflash27b_last_error()); + return false; + } + const int max_ctx = std::max(1, cfg_.device.max_ctx); + if (!create_kimi_k3_cache(backend_, weights_, max_ctx, cache_)) { + std::fprintf(stderr, "[kimi-k3] cache allocation failed (max_ctx=%d)\n", + max_ctx); + return false; + } + std::fprintf(stderr, + "[kimi-k3] native backend ready on device %d (max_ctx=%d, " + "correctness-first sequential prefill)\n", + cfg_.device.primary_gpu(), max_ctx); + std::fflush(stderr); + return true; +} + +void KimiK3Backend::print_ready_banner() const { + std::printf("[kimi-k3-daemon] ready (layers=%d hidden=%d experts=%d " + "vocab=%d max_ctx=%d)\n", + weights_.n_layer, weights_.n_embd, weights_.n_expert, + weights_.n_vocab, cache_.max_ctx); + std::fflush(stdout); +} + +bool KimiK3Backend::park(ParkTarget target) { + if (!park_target_includes_target_model(target)) return false; + if (!parked_) { + free_kimi_k3_weights(weights_); + parked_ = true; + } + return true; +} + +bool KimiK3Backend::unpark(ParkTarget target) { + if (!park_target_includes_target_model(target)) return false; + if (parked_) { + if (!load_kimi_k3_gguf(cfg_.model_path, backend_, weights_)) return false; + parked_ = false; + } + return true; +} + +int32_t KimiK3Backend::choose_token(const std::vector & logits, + const SamplerCfg & sampler, + const std::vector & history) { + if (sampler.needs_logit_processing()) { + return sample_logits(logits.data(), weights_.n_vocab, + sampler, history, rng_); + } + return static_cast(std::distance(logits.begin(), + std::max_element(logits.begin(), logits.end()))); +} + +GenerateResult KimiK3Backend::generate_impl(const GenerateRequest & req, + const DaemonIO & io) { + GenerateResult result; + DaemonIO out_io = io.with_token_callback(req.on_token); + if (parked_) { + result.fail(GenerateErrorCode::ModelParked); + out_io.emit(-1); + return result; + } + if (req.prompt.empty()) { + result.fail(GenerateErrorCode::PrefillFailed, "empty prompt"); + out_io.emit(-1); + return result; + } + if (req.prompt.size() + static_cast(std::max(0, req.n_gen)) > + static_cast(cache_.max_ctx)) { + result.fail(GenerateErrorCode::ContextOverflow, + "prompt plus generation exceeds Kimi-K3 cache"); + out_io.emit(-1); + return result; + } + if (req.do_sample && req.sampler.seed != 0) rng_.seed(req.sampler.seed); + + reset_kimi_k3_cache(cache_); + std::vector logits; + const auto prefill_begin = std::chrono::steady_clock::now(); + for (size_t i = 0; i < req.prompt.size(); ++i) { + if (!kimi_k3_step(backend_, weights_, cache_, req.prompt[i], + static_cast(i), logits)) { + result.fail(GenerateErrorCode::PrefillFailed, + dflash27b_last_error()); + out_io.emit(-1); + return result; + } + } + const auto prefill_end = std::chrono::steady_clock::now(); + result.prefill_s = std::chrono::duration(prefill_end - prefill_begin).count(); + + const auto decode_begin = std::chrono::steady_clock::now(); + bool budget_close_started = false; + size_t close_inject_pos = 0; + for (int i = 0; i < req.n_gen; ++i) { + int32_t next = choose_token(logits, req.sampler, result.tokens); + + // Preserve the shared Level-2 budget contract even before speculative + // decode support lands for Kimi-K3. + const auto & close_ids = req.budget_hook.close_token_ids; + if (!close_ids.empty()) { + if (budget_close_started && close_inject_pos < close_ids.size()) { + next = close_ids[close_inject_pos++]; + result.budget_forced_close = true; + } else if (!budget_close_started && + req.n_gen - i <= + req.budget_hook.hard_limit_remaining) { + budget_close_started = true; + if (next == close_ids.front()) { + close_inject_pos = 1; + } else { + next = close_ids.front(); + close_inject_pos = 1; + result.budget_forced_close = true; + } + } + } + + result.tokens.push_back(next); + out_io.emit(next); + if (out_io.cancelled || next == weights_.eos_token_id) break; + if (i + 1 < req.n_gen) { + if (!kimi_k3_step(backend_, weights_, cache_, next, + cache_.cur_pos, logits)) { + result.fail(GenerateErrorCode::DecodeFailed, + dflash27b_last_error()); + out_io.emit(-1); + return result; + } + } + } + const auto decode_end = std::chrono::steady_clock::now(); + result.decode_s = std::chrono::duration(decode_end - decode_begin).count(); + out_io.emit(-1); + result.succeed(); + return result; +} + +bool KimiK3Backend::snapshot_save(int slot) { + (void)slot; + return false; +} + +void KimiK3Backend::snapshot_free(int slot) { + (void)slot; +} + +bool KimiK3Backend::snapshot_used(int slot) const { + (void)slot; + return false; +} + +int KimiK3Backend::snapshot_cur_pos(int slot) const { + (void)slot; + return 0; +} + +GenerateResult KimiK3Backend::restore_and_generate_impl( + int slot, const GenerateRequest & req, const DaemonIO & io) { + (void)slot; + (void)req; + GenerateResult result; + result.fail(GenerateErrorCode::InvalidSnapshotSlot, + "Kimi-K3 prefix snapshots are not implemented yet"); + io.emit(-1); + return result; +} + +bool KimiK3Backend::handle_compress(const std::string & line, + const DaemonIO & io) { + (void)line; + (void)io; + return false; +} + +void KimiK3Backend::shutdown() { + free_kimi_k3_cache(cache_); + free_kimi_k3_weights(weights_); + if (backend_) { + ggml_backend_free(backend_); + backend_ = nullptr; + } + parked_ = false; +} + +} // namespace dflash::common diff --git a/server/src/kimi_k3/kimi_k3_backend.h b/server/src/kimi_k3/kimi_k3_backend.h new file mode 100644 index 000000000..d15d37416 --- /dev/null +++ b/server/src/kimi_k3/kimi_k3_backend.h @@ -0,0 +1,59 @@ +#pragma once + +#include "common/model_backend.h" +#include "kimi_k3_internal.h" +#include "placement/placement_config.h" + +#include +#include + +namespace dflash::common { + +struct KimiK3BackendConfig { + const char * model_path = nullptr; + DevicePlacement device; + int stream_fd = -1; +}; + +class KimiK3Backend final : public ModelBackend { +public: + explicit KimiK3Backend(const KimiK3BackendConfig & cfg); + ~KimiK3Backend() override; + + bool init(); + + void print_ready_banner() const override; + bool park(ParkTarget target) override; + bool unpark(ParkTarget target) override; + bool is_target_parked() const override { return parked_; } + + GenerateResult generate_impl(const GenerateRequest & req, + const DaemonIO & io) override; + GenerateResult restore_and_generate_impl(int slot, + const GenerateRequest & req, + const DaemonIO & io) override; + + bool snapshot_save(int slot) override; + void snapshot_free(int slot) override; + bool snapshot_used(int slot) const override; + int snapshot_cur_pos(int slot) const override; + + bool handle_compress(const std::string & line, + const DaemonIO & io) override; + void free_drafter() override {} + void shutdown() override; + +private: + int32_t choose_token(const std::vector & logits, + const SamplerCfg & sampler, + const std::vector & history); + + KimiK3BackendConfig cfg_; + ggml_backend_t backend_ = nullptr; + KimiK3Weights weights_; + KimiK3Cache cache_; + bool parked_ = false; + std::mt19937_64 rng_{std::random_device{}()}; +}; + +} // namespace dflash::common diff --git a/server/src/kimi_k3/kimi_k3_graph.cpp b/server/src/kimi_k3/kimi_k3_graph.cpp new file mode 100644 index 000000000..b5e601447 --- /dev/null +++ b/server/src/kimi_k3/kimi_k3_graph.cpp @@ -0,0 +1,496 @@ +#include "kimi_k3_internal.h" + +#include "common/moe_router_graph.h" +#include "internal.h" + +#include "ggml-alloc.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { +namespace { + +ggml_tensor * rms_norm(ggml_context * ctx, + ggml_tensor * x, + ggml_tensor * weight, + float eps) { + x = ggml_rms_norm(ctx, x, eps); + return weight ? ggml_mul(ctx, x, weight) : x; +} + +ggml_tensor * situ(ggml_context * ctx, + ggml_tensor * gate, + ggml_tensor * up, + float beta, + float linear_beta) { + ggml_tensor * a = ggml_scale(ctx, + ggml_tanh(ctx, ggml_scale(ctx, gate, 1.0f / beta)), beta); + a = ggml_mul(ctx, a, ggml_sigmoid(ctx, gate)); + if (linear_beta > 0.0f) { + up = ggml_scale(ctx, + ggml_tanh(ctx, ggml_scale(ctx, up, 1.0f / linear_beta)), + linear_beta); + } + return ggml_mul(ctx, a, up); +} + +struct AttnResBank { + ggml_context * ctx = nullptr; + float eps = 1.0e-5f; + int64_t n_embd = 0; + std::vector checkpoints; + ggml_tensor * stack = nullptr; + size_t stack_size = 0; + + void push(ggml_tensor * cur) { + checkpoints.push_back(ggml_reshape_3d(ctx, cur, n_embd, 1, 1)); + } + + ggml_tensor * get_stack() { + if (stack && stack_size == checkpoints.size()) return stack; + stack = checkpoints.front(); + for (size_t i = 1; i < checkpoints.size(); ++i) { + stack = ggml_concat(ctx, stack, checkpoints[i], 1); + } + stack_size = checkpoints.size(); + return stack; + } + + ggml_tensor * mix(ggml_tensor * cur, ggml_tensor * score_weight) { + if (checkpoints.empty()) return cur; + const int64_t n = static_cast(checkpoints.size()); + ggml_tensor * src = get_stack(); // [hidden, n_checkpoint, 1] + + ggml_tensor * score_src = rms_norm(ctx, src, score_weight, eps); + score_src = ggml_sum_rows(ctx, score_src); + score_src = ggml_reshape_2d(ctx, score_src, n, 1); + + ggml_tensor * score_cur = rms_norm(ctx, cur, score_weight, eps); + score_cur = ggml_sum_rows(ctx, score_cur); + + ggml_tensor * probs = ggml_soft_max(ctx, + ggml_concat(ctx, score_src, score_cur, 0)); + ggml_tensor * p_src = ggml_cont(ctx, + ggml_view_2d(ctx, probs, n, 1, probs->nb[1], 0)); + ggml_tensor * p_cur = ggml_cont(ctx, + ggml_view_2d(ctx, probs, 1, 1, probs->nb[1], + probs->nb[0] * static_cast(n))); + + // Reduce checkpoint dimension with an ordinary matrix product. The + // newer upstream ggml has a dedicated dsv4_hc_pre op for this exact + // contraction; the Lucebox ggml snapshot predates that API, and this + // algebra is identical: [checkpoint, hidden] x [checkpoint, 1]. + ggml_tensor * src_t = ggml_cont(ctx, + ggml_permute(ctx, src, 1, 0, 2, 3)); + ggml_tensor * out = ggml_mul_mat(ctx, src_t, p_src); + return ggml_add(ctx, out, ggml_mul(ctx, cur, p_cur)); + } +}; + +ggml_tensor * kda_conv1d(ggml_context * ctx, + ggml_cgraph * graph, + ggml_tensor * all_state, + int qkv, + ggml_tensor * x, + ggml_tensor * projection, + ggml_tensor * conv_weight, + int d_conv, + int head_dim, + int n_head) { + const int64_t d_inner = static_cast(head_dim) * n_head; + const int64_t state_rows = d_conv - 1; + const size_t block_offset = static_cast(qkv) * d_inner * all_state->nb[1]; + ggml_tensor * state = ggml_view_3d(ctx, all_state, + state_rows, d_inner, 1, all_state->nb[1], all_state->nb[2], + block_offset); + + ggml_tensor * projected = ggml_mul_mat(ctx, projection, x); + projected = ggml_reshape_3d(ctx, projected, d_inner, 1, 1); + ggml_tensor * conv_input = ggml_concat(ctx, state, + ggml_transpose(ctx, projected), 0); + + // Drop the oldest row and persist the newest d_conv-1 values. + ggml_tensor * newest = ggml_view_3d(ctx, conv_input, + state_rows, d_inner, 1, conv_input->nb[1], conv_input->nb[2], + conv_input->nb[0]); + ggml_build_forward_expand(graph, ggml_cpy(ctx, newest, state)); + + ggml_tensor * cw = ggml_reshape_2d(ctx, conv_weight, d_conv, d_inner); + ggml_tensor * out = ggml_silu(ctx, ggml_ssm_conv(ctx, conv_input, cw)); + out = ggml_reshape_4d(ctx, out, head_dim, n_head, 1, 1); + return out; +} + +ggml_tensor * build_kda(ggml_context * ctx, + ggml_cgraph * graph, + const KimiK3Weights & w, + const KimiK3Layer & layer, + KimiK3LayerCache & cache, + ggml_tensor * cur) { + const int head_dim = w.kda_head_dim; + const int n_head = w.n_head; + const int64_t d_inner = static_cast(head_dim) * n_head; + + ggml_tensor * q = kda_conv1d(ctx, graph, cache.conv_state, 0, cur, + layer.wq, layer.ssm_q_conv, w.ssm_d_conv, head_dim, n_head); + ggml_tensor * k = kda_conv1d(ctx, graph, cache.conv_state, 1, cur, + layer.wk, layer.ssm_k_conv, w.ssm_d_conv, head_dim, n_head); + ggml_tensor * v = kda_conv1d(ctx, graph, cache.conv_state, 2, cur, + layer.wv, layer.ssm_v_conv, w.ssm_d_conv, head_dim, n_head); + + ggml_tensor * decay = ggml_mul_mat(ctx, layer.ssm_f_a, cur); + decay = ggml_mul_mat(ctx, layer.ssm_f_b, decay); + decay = ggml_add(ctx, decay, layer.ssm_dt_b); + ggml_tensor * A = ggml_reshape_3d(ctx, layer.ssm_a, 1, n_head, 1); + if (std::isfinite(w.kda_gate_lower_bound)) { + decay = ggml_reshape_3d(ctx, decay, head_dim, n_head, 1); + decay = ggml_mul(ctx, decay, A); + decay = ggml_sigmoid(ctx, ggml_scale(ctx, decay, -1.0f)); + decay = ggml_scale(ctx, decay, w.kda_gate_lower_bound); + } else { + decay = ggml_softplus(ctx, decay); + decay = ggml_reshape_3d(ctx, decay, head_dim, n_head, 1); + decay = ggml_mul(ctx, decay, A); + } + decay = ggml_reshape_4d(ctx, decay, head_dim, n_head, 1, 1); + + ggml_tensor * beta = ggml_mul_mat(ctx, layer.ssm_beta, cur); + beta = ggml_sigmoid(ctx, ggml_reshape_4d(ctx, beta, 1, n_head, 1, 1)); + + q = ggml_l2_norm(ctx, q, w.rms_eps); + k = ggml_l2_norm(ctx, k, w.rms_eps); + ggml_tensor * state = ggml_reshape_4d(ctx, cache.ssm_state, + head_dim, head_dim, n_head, 1); + ggml_tensor * packed = ggml_gated_delta_net(ctx, q, k, v, decay, beta, state); + ggml_gated_delta_net_set_skip_intermediate(packed, true); + + const size_t elt = ggml_element_size(packed); + ggml_tensor * output = ggml_view_4d(ctx, packed, + head_dim, n_head, 1, 1, + static_cast(head_dim) * elt, + static_cast(head_dim) * n_head * elt, + static_cast(head_dim) * n_head * elt, 0); + ggml_tensor * new_state = ggml_view_4d(ctx, packed, + head_dim, head_dim, n_head, 1, + static_cast(head_dim) * elt, + static_cast(head_dim) * head_dim * elt, + static_cast(head_dim) * head_dim * n_head * elt, + static_cast(head_dim) * n_head * elt); + ggml_build_forward_expand(graph, + ggml_cpy(ctx, new_state, cache.ssm_state)); + + ggml_tensor * gate = ggml_mul_mat(ctx, layer.ssm_g, cur); + gate = ggml_reshape_3d(ctx, gate, head_dim, n_head, 1); + output = ggml_reshape_3d(ctx, output, head_dim, n_head, 1); + output = rms_norm(ctx, output, layer.ssm_o_norm, w.rms_eps); + output = ggml_mul(ctx, output, ggml_sigmoid(ctx, gate)); + output = ggml_cont_2d(ctx, output, d_inner, 1); + return ggml_mul_mat(ctx, layer.wo, output); +} + +ggml_tensor * build_mla(ggml_context * ctx, + ggml_cgraph * graph, + const KimiK3Weights & w, + const KimiK3Layer & layer, + KimiK3LayerCache & cache, + ggml_tensor * cur, + int position) { + const int n_head = w.n_head; + const int kv_rank = w.kv_lora_rank; + const int key_dim = w.mla_k_head_dim; + const int value_dim = w.mla_v_head_dim; + const int rope_dim = w.rope_dim; + const int nope_dim = key_dim - rope_dim; + const int compact_dim = kv_rank + rope_dim; + const int kv_len = position + 1; + + ggml_tensor * gate_input = cur; + ggml_tensor * q_cur = nullptr; + if (layer.wq_a) { + q_cur = ggml_mul_mat(ctx, layer.wq_a, cur); + q_cur = rms_norm(ctx, q_cur, layer.wq_a_norm, w.rms_eps); + q_cur = ggml_mul_mat(ctx, layer.wq_b, q_cur); + } else { + q_cur = ggml_mul_mat(ctx, layer.wq, cur); + } + + ggml_tensor * compact_pe = ggml_mul_mat(ctx, layer.wkv_a_mqa, cur); + ggml_tensor * compact = ggml_view_2d(ctx, compact_pe, kv_rank, 1, + ggml_row_size(compact_pe->type, compact_dim), 0); + ggml_tensor * k_pe = ggml_view_3d(ctx, compact_pe, rope_dim, 1, 1, + ggml_row_size(compact_pe->type, compact_dim), + ggml_row_size(compact_pe->type, compact_dim), + ggml_row_size(compact_pe->type, kv_rank)); + compact = rms_norm(ctx, compact, layer.wkv_a_norm, w.rms_eps); + + ggml_tensor * q_nope = ggml_view_3d(ctx, q_cur, nope_dim, n_head, 1, + ggml_row_size(q_cur->type, key_dim), + ggml_row_size(q_cur->type, key_dim) * n_head, 0); + ggml_tensor * q_pe = ggml_view_3d(ctx, q_cur, rope_dim, n_head, 1, + ggml_row_size(q_cur->type, key_dim), + ggml_row_size(q_cur->type, key_dim) * n_head, + ggml_row_size(q_cur->type, nope_dim)); + q_nope = ggml_permute(ctx, q_nope, 0, 2, 1, 3); + q_nope = ggml_mul_mat(ctx, layer.wk_b, q_nope); + q_nope = ggml_permute(ctx, q_nope, 0, 2, 1, 3); + ggml_tensor * q = ggml_concat(ctx, q_nope, q_pe, 0); + + ggml_tensor * compact_3d = ggml_reshape_3d(ctx, compact, kv_rank, 1, 1); + ggml_tensor * current_k = ggml_concat(ctx, compact_3d, k_pe, 0); + + ggml_tensor * dst = ggml_view_3d(ctx, cache.mla_k, + compact_dim, 1, 1, cache.mla_k->nb[1], cache.mla_k->nb[2], + static_cast(position) * cache.mla_k->nb[2]); + ggml_build_forward_expand(graph, ggml_cpy(ctx, current_k, dst)); + + ggml_tensor * k = ggml_view_3d(ctx, cache.mla_k, + compact_dim, 1, kv_len, cache.mla_k->nb[1], cache.mla_k->nb[2], 0); + ggml_tensor * v = ggml_view_3d(ctx, k, + kv_rank, 1, kv_len, k->nb[1], k->nb[2], 0); + + // Same non-flash absorbed-MLA algebra as llama.cpp. Avoiding flash here + // is intentional: current upstream K3 support also disables it for this + // graph, and this path is the numerical oracle for later fused kernels. + const bool v_trans = v->nb[1] > v->nb[2]; + q = ggml_permute(ctx, q, 0, 2, 1, 3); + k = ggml_permute(ctx, k, 0, 2, 1, 3); + v = ggml_permute(ctx, v, 0, 2, 1, 3); + ggml_tensor * scores = ggml_mul_mat(ctx, k, q); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + scores = ggml_soft_max_ext(ctx, scores, nullptr, + 1.0f / std::sqrt(static_cast(key_dim)), + 0.0f); + if (!v_trans) v = ggml_cont(ctx, ggml_transpose(ctx, v)); + ggml_tensor * out = ggml_mul_mat(ctx, v, scores); + out = ggml_mul_mat(ctx, layer.wv_b, out); + out = ggml_permute(ctx, out, 0, 2, 1, 3); + out = ggml_cont_2d(ctx, out, + static_cast(value_dim) * n_head, 1); + + if (layer.wqkv_gate) { + ggml_tensor * output_gate = ggml_sigmoid(ctx, + ggml_mul_mat(ctx, layer.wqkv_gate, gate_input)); + out = ggml_mul(ctx, out, output_gate); + } + return ggml_mul_mat(ctx, layer.wo, out); +} + +ggml_tensor * build_latent_moe(ggml_context * ctx, + ggml_cgraph * graph, + const KimiK3Weights & w, + const KimiK3Layer & layer, + ggml_tensor * cur) { + ggml_tensor * identity = cur; + ggml_tensor * routed_in = ggml_mul_mat(ctx, layer.ffn_routed_down, cur); + ggml_tensor * logits = ggml_mul_mat(ctx, layer.ffn_gate_inp, identity); + + TopKMoeRouterResult router; + if (w.expert_gating_func == 2) { + router = build_sigmoid_topk_moe_router(ctx, graph, logits, + layer.ffn_exp_probs_b, w.n_expert, w.n_expert_used, 1, + w.expert_weights_norm, w.expert_weights_scale, false); + } else { + ggml_tensor * probs = ggml_soft_max(ctx, logits); + ggml_tensor * selected = ggml_argsort_top_k(ctx, probs, w.n_expert_used); + ggml_tensor * probs_3d = ggml_reshape_3d(ctx, probs, 1, w.n_expert, 1); + ggml_tensor * weights = ggml_get_rows(ctx, probs_3d, selected); + weights = ggml_reshape_2d(ctx, weights, w.n_expert_used, 1); + if (w.expert_weights_norm) { + ggml_tensor * sum = ggml_clamp(ctx, ggml_sum_rows(ctx, weights), + 6.103515625e-5f, INFINITY); + weights = ggml_div(ctx, weights, sum); + } + if (w.expert_weights_scale != 1.0f) { + weights = ggml_scale(ctx, weights, w.expert_weights_scale); + } + router.selected = selected; + router.weights_2d = weights; + router.weights_3d = ggml_reshape_3d(ctx, weights, 1, w.n_expert_used, 1); + } + + ggml_tensor * routed_3d = ggml_reshape_3d(ctx, routed_in, + w.n_expert_latent, 1, 1); + ggml_tensor * gate = ggml_mul_mat_id(ctx, layer.ffn_gate_exps, + routed_3d, router.selected); + ggml_tensor * up = ggml_mul_mat_id(ctx, layer.ffn_up_exps, + routed_3d, router.selected); + ggml_tensor * activated = situ(ctx, gate, up, + w.situ_beta, w.situ_linear_beta); + ggml_tensor * experts = ggml_mul_mat_id(ctx, layer.ffn_down_exps, + activated, router.selected); + experts = ggml_mul(ctx, experts, router.weights_3d); + ggml_tensor * sum_shape = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, + w.n_expert_latent, 1, 1); + ggml_tensor * moe = ggml_repeat_back(ctx, experts, sum_shape); + moe = ggml_reshape_2d(ctx, moe, w.n_expert_latent, 1); + if (layer.ffn_routed_norm) { + moe = rms_norm(ctx, moe, layer.ffn_routed_norm, w.rms_eps); + } + moe = ggml_mul_mat(ctx, layer.ffn_routed_up, moe); + + ggml_tensor * shared_gate = ggml_mul_mat(ctx, layer.ffn_gate_shexp, identity); + ggml_tensor * shared_up = ggml_mul_mat(ctx, layer.ffn_up_shexp, identity); + ggml_tensor * shared = situ(ctx, shared_gate, shared_up, + w.situ_beta, w.situ_linear_beta); + shared = ggml_mul_mat(ctx, layer.ffn_down_shexp, shared); + return ggml_add(ctx, moe, shared); +} + +} // namespace + +bool create_kimi_k3_cache(ggml_backend_t backend, + const KimiK3Weights & w, + int max_ctx, + KimiK3Cache & out) { + free_kimi_k3_cache(out); + if (!backend || max_ctx <= 0) return false; + + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * + static_cast(w.n_layer * 3 + 16) + 16384; + params.no_alloc = true; + out.ctx = ggml_init(params); + if (!out.ctx) return false; + + out.layers.resize(static_cast(w.n_layer)); + const int64_t d_inner = static_cast(w.kda_head_dim) * w.n_head; + const int compact_dim = w.kv_lora_rank + w.rope_dim; + for (int il = 0; il < w.n_layer; ++il) { + KimiK3LayerCache & layer_cache = out.layers[static_cast(il)]; + char name[80]; + if (w.layers[static_cast(il)].recurrent) { + layer_cache.conv_state = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, + w.ssm_d_conv - 1, 3 * d_inner); + layer_cache.ssm_state = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F32, + w.kda_head_dim, w.kda_head_dim, w.n_head); + std::snprintf(name, sizeof(name), "kimi_k3_conv_state_%d", il); + ggml_set_name(layer_cache.conv_state, name); + std::snprintf(name, sizeof(name), "kimi_k3_ssm_state_%d", il); + ggml_set_name(layer_cache.ssm_state, name); + } else { + layer_cache.mla_k = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F16, + compact_dim, 1, max_ctx); + std::snprintf(name, sizeof(name), "kimi_k3_mla_k_%d", il); + ggml_set_name(layer_cache.mla_k, name); + } + } + + out.buf = ggml_backend_alloc_ctx_tensors(out.ctx, backend); + if (!out.buf) { + free_kimi_k3_cache(out); + return false; + } + out.max_ctx = max_ctx; + reset_kimi_k3_cache(out); + return true; +} + +void reset_kimi_k3_cache(KimiK3Cache & cache) { + if (cache.buf) ggml_backend_buffer_clear(cache.buf, 0); + cache.cur_pos = 0; +} + +void free_kimi_k3_cache(KimiK3Cache & cache) { + if (cache.buf) ggml_backend_buffer_free(cache.buf); + if (cache.ctx) ggml_free(cache.ctx); + cache = KimiK3Cache{}; +} + +bool kimi_k3_step(ggml_backend_t backend, + const KimiK3Weights & w, + KimiK3Cache & cache, + int32_t token, + int position, + std::vector & logits) { + if (!backend || !w.ctx || !cache.ctx || position < 0 || + position >= cache.max_ctx || position != cache.cur_pos || + token < 0 || token >= w.n_vocab) { + set_last_error("Kimi-K3 step: invalid backend, cache position, or token"); + return false; + } + + ggml_init_params params{}; + params.mem_size = 64ull * 1024ull * 1024ull; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + set_last_error("Kimi-K3 step: graph context allocation failed"); + return false; + } + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 32768, false); + + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_set_name(ids, "token_id"); + ggml_set_input(ids); + ggml_tensor * hidden = ggml_get_rows(ctx, w.tok_embd, ids); + + AttnResBank residuals; + residuals.ctx = ctx; + residuals.eps = w.rms_eps; + residuals.n_embd = w.n_embd; + for (int il = 0; il < w.n_layer; ++il) { + const KimiK3Layer & layer = w.layers[static_cast(il)]; + KimiK3LayerCache & layer_cache = cache.layers[static_cast(il)]; + ggml_tensor * prefix = hidden; + ggml_tensor * cur = residuals.mix(prefix, layer.attn_res_score); + const bool banked = il % w.attn_res_block_size == 0; + if (banked) residuals.push(prefix); + + cur = rms_norm(ctx, cur, layer.attn_norm, w.rms_eps); + cur = layer.recurrent + ? build_kda(ctx, graph, w, layer, layer_cache, cur) + : build_mla(ctx, graph, w, layer, layer_cache, cur, position); + prefix = banked ? cur : ggml_add(ctx, prefix, cur); + + cur = residuals.mix(prefix, layer.ffn_res_score); + cur = rms_norm(ctx, cur, layer.ffn_norm, w.rms_eps); + if (il < w.n_dense_lead) { + ggml_tensor * gate = ggml_mul_mat(ctx, layer.ffn_gate, cur); + ggml_tensor * up = ggml_mul_mat(ctx, layer.ffn_up, cur); + cur = situ(ctx, gate, up, w.situ_beta, w.situ_linear_beta); + cur = ggml_mul_mat(ctx, layer.ffn_down, cur); + } else { + cur = build_latent_moe(ctx, graph, w, layer, cur); + } + hidden = ggml_add(ctx, prefix, cur); + } + + hidden = residuals.mix(hidden, w.output_res_score); + hidden = rms_norm(ctx, hidden, w.output_norm, w.rms_eps); + ggml_tensor * output = ggml_mul_mat(ctx, w.output, hidden); + ggml_set_name(output, "logits"); + ggml_set_output(output); + ggml_build_forward_expand(graph, output); + + ggml_gallocr_t allocator = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + if (!allocator || !ggml_gallocr_alloc_graph(allocator, graph)) { + set_last_error("Kimi-K3 step: graph allocation failed"); + if (allocator) ggml_gallocr_free(allocator); + ggml_free(ctx); + return false; + } + ggml_backend_tensor_set(ids, &token, 0, sizeof(token)); + const ggml_status status = ggml_backend_graph_compute(backend, graph); + if (status != GGML_STATUS_SUCCESS) { + set_last_error("Kimi-K3 step: graph compute failed with status " + + std::to_string(static_cast(status))); + ggml_gallocr_free(allocator); + ggml_free(ctx); + return false; + } + + logits.resize(static_cast(w.n_vocab)); + ggml_backend_tensor_get(output, logits.data(), 0, + logits.size() * sizeof(float)); + cache.cur_pos = position + 1; + ggml_gallocr_free(allocator); + ggml_free(ctx); + return true; +} + +} // namespace dflash::common diff --git a/server/src/kimi_k3/kimi_k3_internal.h b/server/src/kimi_k3/kimi_k3_internal.h new file mode 100644 index 000000000..a1296bade --- /dev/null +++ b/server/src/kimi_k3/kimi_k3_internal.h @@ -0,0 +1,159 @@ +// Native Kimi-K3 text-model support for Lucebox Hub. +// +// This is intentionally split into three model-neutral boundaries: +// * GGUF loading owns tensor metadata/storage only; +// * KimiK3Cache owns recurrent/attention state only; and +// * kimi_k3_step owns the architecture graph only. +// +// Routed-expert placement can therefore replace the resident expert tensors +// with the common MoE stream engine without changing KDA, MLA, AttnRes, or the +// public ModelBackend contract. + +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include + +namespace dflash::common { + +struct KimiK3Layer { + bool recurrent = false; + + ggml_tensor * attn_norm = nullptr; + ggml_tensor * ffn_norm = nullptr; + ggml_tensor * attn_res_score = nullptr; + ggml_tensor * ffn_res_score = nullptr; + + // KDA (recurrent) attention. + ggml_tensor * wq = nullptr; + ggml_tensor * wk = nullptr; + ggml_tensor * wv = nullptr; + ggml_tensor * wo = nullptr; + ggml_tensor * ssm_q_conv = nullptr; + ggml_tensor * ssm_k_conv = nullptr; + ggml_tensor * ssm_v_conv = nullptr; + ggml_tensor * ssm_f_a = nullptr; + ggml_tensor * ssm_f_b = nullptr; + ggml_tensor * ssm_beta = nullptr; + ggml_tensor * ssm_a = nullptr; + ggml_tensor * ssm_dt_b = nullptr; + ggml_tensor * ssm_g = nullptr; + ggml_tensor * ssm_o_norm = nullptr; + + // MLA attention. Kimi-K3 uses the absorbed K-only cache when wk_b/wv_b + // are present, which is the layout emitted by the official converter. + ggml_tensor * wq_a = nullptr; + ggml_tensor * wq_a_norm = nullptr; + ggml_tensor * wq_b = nullptr; + ggml_tensor * wkv_a_mqa = nullptr; + ggml_tensor * wkv_a_norm = nullptr; + ggml_tensor * wk_b = nullptr; + ggml_tensor * wv_b = nullptr; + ggml_tensor * wkv_b = nullptr; + ggml_tensor * wqkv_gate = nullptr; + + // Dense FFN (leading dense blocks). + ggml_tensor * ffn_gate = nullptr; + ggml_tensor * ffn_up = nullptr; + ggml_tensor * ffn_down = nullptr; + + // Latent routed MoE + full-width shared expert. + ggml_tensor * ffn_gate_inp = nullptr; + ggml_tensor * ffn_exp_probs_b = nullptr; + ggml_tensor * ffn_gate_exps = nullptr; + ggml_tensor * ffn_up_exps = nullptr; + ggml_tensor * ffn_down_exps = nullptr; + ggml_tensor * ffn_routed_down = nullptr; + ggml_tensor * ffn_routed_up = nullptr; + ggml_tensor * ffn_routed_norm = nullptr; + ggml_tensor * ffn_gate_shexp = nullptr; + ggml_tensor * ffn_up_shexp = nullptr; + ggml_tensor * ffn_down_shexp = nullptr; +}; + +struct KimiK3Weights { + ggml_context * ctx = nullptr; + ggml_backend_t backend = nullptr; + ggml_backend_buffer_t buf = nullptr; + + ggml_tensor * tok_embd = nullptr; + ggml_tensor * output_norm = nullptr; + ggml_tensor * output = nullptr; + ggml_tensor * output_res_score = nullptr; + std::vector layers; + + int n_layer = 0; + int n_embd = 0; + int n_ff = 0; + int n_vocab = 0; + int n_ctx_train = 0; + int n_head = 0; + int n_expert = 0; + int n_expert_used = 0; + int n_ff_exp = 0; + int n_expert_latent = 0; + int n_expert_shared = 0; + int n_dense_lead = 0; + + int ssm_d_conv = 0; + int kda_head_dim = 0; + int q_lora_rank = 0; + int kv_lora_rank = 0; + int mla_k_head_dim = 0; + int mla_v_head_dim = 0; + int rope_dim = 0; + int attn_res_block_size = 0; + + float rms_eps = 1.0e-5f; + float kda_gate_lower_bound = -INFINITY; + float expert_weights_scale = 1.0f; + bool expert_weights_norm = true; + int expert_gating_func = 2; // sigmoid + float situ_beta = 4.0f; + float situ_linear_beta = 25.0f; + int32_t eos_token_id = 2; +}; + +struct KimiK3LayerCache { + ggml_tensor * conv_state = nullptr; // [d_conv-1, 3*d_inner], F32 + ggml_tensor * ssm_state = nullptr; // [head_dim, head_dim, n_head], F32 + ggml_tensor * mla_k = nullptr; // [kv_rank+rope_dim, 1, max_ctx], F16 +}; + +struct KimiK3Cache { + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buf = nullptr; + std::vector layers; + int max_ctx = 0; + int cur_pos = 0; +}; + +bool load_kimi_k3_gguf(const std::string & path, + ggml_backend_t backend, + KimiK3Weights & out); +void free_kimi_k3_weights(KimiK3Weights & w); + +bool create_kimi_k3_cache(ggml_backend_t backend, + const KimiK3Weights & w, + int max_ctx, + KimiK3Cache & out); +void reset_kimi_k3_cache(KimiK3Cache & cache); +void free_kimi_k3_cache(KimiK3Cache & cache); + +// Executes exactly one token. Token-at-a-time is deliberate for the first +// correctness path: it makes the recurrent state transition explicit and is +// numerically equivalent to chunked prefill. Persistent/captured decode graphs +// and chunked KDA are performance layers added above this contract. +bool kimi_k3_step(ggml_backend_t backend, + const KimiK3Weights & w, + KimiK3Cache & cache, + int32_t token, + int position, + std::vector & logits); + +} // namespace dflash::common diff --git a/server/src/kimi_k3/kimi_k3_loader.cpp b/server/src/kimi_k3/kimi_k3_loader.cpp new file mode 100644 index 000000000..9aa8a92f3 --- /dev/null +++ b/server/src/kimi_k3/kimi_k3_loader.cpp @@ -0,0 +1,326 @@ +#include "kimi_k3_internal.h" + +#include "common/gguf_bounds.h" +#include "common/gguf_mmap.h" +#include "internal.h" + +#include +#include +#include +#include +#include +#include + +namespace dflash::common { +namespace { + +uint32_t get_u32_or(const gguf_context * g, const char * key, uint32_t fallback) { + const int64_t id = gguf_find_key(g, key); + if (id < 0) return fallback; + if (gguf_get_kv_type(g, id) == GGUF_TYPE_ARRAY) { + if (gguf_get_arr_n(g, id) == 0) return fallback; + const gguf_type type = gguf_get_arr_type(g, id); + const void * data = gguf_get_arr_data(g, id); + if (type == GGUF_TYPE_UINT32) return static_cast(data)[0]; + if (type == GGUF_TYPE_INT32) return static_cast(static_cast(data)[0]); + return fallback; + } + const gguf_type type = gguf_get_kv_type(g, id); + if (type == GGUF_TYPE_UINT32) return gguf_get_val_u32(g, id); + if (type == GGUF_TYPE_INT32) return static_cast(gguf_get_val_i32(g, id)); + return fallback; +} + +float get_f32_or(const gguf_context * g, const char * key, float fallback) { + const int64_t id = gguf_find_key(g, key); + if (id < 0) return fallback; + if (gguf_get_kv_type(g, id) == GGUF_TYPE_FLOAT32) return gguf_get_val_f32(g, id); + return fallback; +} + +bool get_bool_or(const gguf_context * g, const char * key, bool fallback) { + const int64_t id = gguf_find_key(g, key); + if (id < 0 || gguf_get_kv_type(g, id) != GGUF_TYPE_BOOL) return fallback; + return gguf_get_val_bool(g, id); +} + +std::vector get_u32_array(const gguf_context * g, const char * key) { + std::vector out; + const int64_t id = gguf_find_key(g, key); + if (id < 0 || gguf_get_kv_type(g, id) != GGUF_TYPE_ARRAY) return out; + const size_t n = gguf_get_arr_n(g, id); + const void * data = gguf_get_arr_data(g, id); + if (gguf_get_arr_type(g, id) == GGUF_TYPE_UINT32) { + const auto * p = static_cast(data); + out.assign(p, p + n); + } else if (gguf_get_arr_type(g, id) == GGUF_TYPE_INT32) { + const auto * p = static_cast(data); + out.reserve(n); + for (size_t i = 0; i < n; ++i) out.push_back(static_cast(p[i])); + } + return out; +} + +bool tensor_shape_is(const ggml_tensor * t, + int64_t ne0, + int64_t ne1 = 1, + int64_t ne2 = 1) { + return t && t->ne[0] == ne0 && t->ne[1] == ne1 && t->ne[2] == ne2; +} + +} // namespace + +bool load_kimi_k3_gguf(const std::string & path, + ggml_backend_t backend, + KimiK3Weights & out) { + free_kimi_k3_weights(out); + + ggml_context * meta_ctx = nullptr; + gguf_init_params params{}; + params.no_alloc = true; + params.ctx = &meta_ctx; + gguf_context * gctx = gguf_init_from_file(path.c_str(), params); + if (!gctx || !meta_ctx) { + set_last_error("Kimi-K3: failed to parse GGUF: " + path); + if (gctx) gguf_free(gctx); + if (meta_ctx) ggml_free(meta_ctx); + return false; + } + + auto fail = [&](const std::string & message) { + set_last_error("Kimi-K3: " + message); + if (out.buf) { + ggml_backend_buffer_free(out.buf); + out.buf = nullptr; + } + gguf_free(gctx); + ggml_free(meta_ctx); + out = KimiK3Weights{}; + return false; + }; + + const int64_t arch_id = gguf_find_key(gctx, "general.architecture"); + if (arch_id < 0 || std::strcmp(gguf_get_val_str(gctx, arch_id), "kimi-k3") != 0) { + return fail("general.architecture must be kimi-k3"); + } + + constexpr const char * A = "kimi-k3."; + auto key = [&](const char * suffix) { return std::string(A) + suffix; }; + auto u32 = [&](const char * suffix, uint32_t fallback = 0) { + const std::string k = key(suffix); + return get_u32_or(gctx, k.c_str(), fallback); + }; + auto f32 = [&](const char * suffix, float fallback) { + const std::string k = key(suffix); + return get_f32_or(gctx, k.c_str(), fallback); + }; + auto boolean = [&](const char * suffix, bool fallback) { + const std::string k = key(suffix); + return get_bool_or(gctx, k.c_str(), fallback); + }; + + out.ctx = meta_ctx; + out.backend = backend; + out.n_layer = static_cast(u32("block_count")); + out.n_embd = static_cast(u32("embedding_length")); + out.n_ff = static_cast(u32("feed_forward_length")); + out.n_vocab = static_cast(u32("vocab_size")); + out.n_ctx_train = static_cast(u32("context_length")); + out.n_head = static_cast(u32("attention.head_count")); + out.n_expert = static_cast(u32("expert_count")); + out.n_expert_used = static_cast(u32("expert_used_count")); + out.n_ff_exp = static_cast(u32("expert_feed_forward_length")); + out.n_expert_latent = static_cast(u32("expert_latent_length")); + out.n_expert_shared = static_cast(u32("expert_shared_count", 1)); + out.n_dense_lead = static_cast(u32("leading_dense_block_count", 0)); + out.ssm_d_conv = static_cast(u32("ssm.conv_kernel")); + out.kda_head_dim = static_cast(u32("kda.head_dim")); + out.q_lora_rank = static_cast(u32("attention.q_lora_rank")); + out.kv_lora_rank = static_cast(u32("attention.kv_lora_rank")); + out.mla_k_head_dim = static_cast(u32("attention.key_length_mla")); + out.mla_v_head_dim = static_cast(u32("attention.value_length_mla")); + out.rope_dim = static_cast(u32("rope.dimension_count")); + out.attn_res_block_size = static_cast(u32("attn_res.block_size")); + out.rms_eps = f32("attention.layer_norm_rms_epsilon", 1.0e-5f); + out.kda_gate_lower_bound = f32("kda.gate_lower_bound", -INFINITY); + out.expert_weights_scale = f32("expert_weights_scale", 1.0f); + out.expert_weights_norm = boolean("expert_weights_norm", true); + out.expert_gating_func = static_cast(u32("expert_gating_func", 2)); + out.situ_beta = f32("activation.situ_beta", 4.0f); + out.situ_linear_beta = f32("activation.situ_linear_beta", 25.0f); + out.eos_token_id = static_cast(get_u32_or(gctx, "tokenizer.ggml.eos_token_id", 2)); + + auto get = [&](const char * name) { return ggml_get_tensor(meta_ctx, name); }; + out.tok_embd = get("token_embd.weight"); + out.output_norm = get("output_norm.weight"); + out.output = get("output.weight"); + out.output_res_score = get("output_res_score.weight"); + if (out.n_vocab == 0 && out.tok_embd) out.n_vocab = static_cast(out.tok_embd->ne[1]); + + constexpr int MAX_LAYERS = 1024; + constexpr int MAX_HEADS = 1024; + constexpr int MAX_EXPERTS = 4096; + if (out.n_layer <= 0 || out.n_layer > MAX_LAYERS || + out.n_embd <= 0 || out.n_head <= 0 || out.n_head > MAX_HEADS || + out.n_vocab <= 0 || out.n_expert <= 0 || out.n_expert > MAX_EXPERTS || + out.n_expert_used <= 0 || out.n_expert_used > out.n_expert || + out.n_expert_latent <= 0 || out.n_ff_exp <= 0 || + out.ssm_d_conv < 2 || out.kda_head_dim <= 0 || + out.attn_res_block_size <= 0 || out.kv_lora_rank <= 0 || + out.mla_k_head_dim <= out.rope_dim || out.mla_v_head_dim <= 0) { + return fail("invalid or incomplete architecture metadata"); + } + if (!tensor_shape_is(out.tok_embd, out.n_embd, out.n_vocab) || + !tensor_shape_is(out.output, out.n_embd, out.n_vocab) || + !tensor_shape_is(out.output_norm, out.n_embd) || + !tensor_shape_is(out.output_res_score, out.n_embd)) { + return fail("missing or malformed top-level tensors"); + } + + std::vector head_kv = get_u32_array(gctx, "kimi-k3.attention.head_count_kv"); + if (head_kv.empty()) { + head_kv.assign(static_cast(out.n_layer), + get_u32_or(gctx, "kimi-k3.attention.head_count_kv", 0)); + } + if (head_kv.size() != static_cast(out.n_layer)) { + return fail("attention.head_count_kv must have one value per layer"); + } + + out.layers.assign(static_cast(out.n_layer), KimiK3Layer{}); + for (int il = 0; il < out.n_layer; ++il) { + char name[160]; + auto find = [&](const char * suffix) -> ggml_tensor * { + std::snprintf(name, sizeof(name), "blk.%d.%s", il, suffix); + return get(name); + }; + KimiK3Layer & L = out.layers[static_cast(il)]; + L.recurrent = head_kv[static_cast(il)] == 0; + L.attn_norm = find("attn_norm.weight"); + L.ffn_norm = find("ffn_norm.weight"); + L.attn_res_score = find("attn_res_score.weight"); + L.ffn_res_score = find("ffn_res_score.weight"); + L.wo = find("attn_output.weight"); + if (!L.attn_norm || !L.ffn_norm || !L.attn_res_score || + !L.ffn_res_score || !L.wo) { + return fail("layer " + std::to_string(il) + " is missing common tensors"); + } + + if (L.recurrent) { + L.wq = find("attn_q.weight"); + L.wk = find("attn_k.weight"); + L.wv = find("attn_v.weight"); + L.ssm_q_conv = find("ssm_conv1d_q.weight"); + L.ssm_k_conv = find("ssm_conv1d_k.weight"); + L.ssm_v_conv = find("ssm_conv1d_v.weight"); + L.ssm_f_a = find("ssm_f_a.weight"); + L.ssm_f_b = find("ssm_f_b.weight"); + L.ssm_beta = find("ssm_beta.weight"); + L.ssm_a = find("ssm_a"); + L.ssm_dt_b = find("ssm_dt.bias"); + L.ssm_g = find("ssm_g.weight"); + L.ssm_o_norm = find("ssm_norm.weight"); + if (!L.wq || !L.wk || !L.wv || !L.ssm_q_conv || !L.ssm_k_conv || + !L.ssm_v_conv || !L.ssm_f_a || !L.ssm_f_b || !L.ssm_beta || + !L.ssm_a || !L.ssm_dt_b || !L.ssm_g || !L.ssm_o_norm) { + return fail("KDA layer " + std::to_string(il) + " is incomplete"); + } + } else { + L.wq_a = find("attn_q_a.weight"); + L.wq_a_norm = find("attn_q_a_norm.weight"); + L.wq_b = find("attn_q_b.weight"); + L.wq = find("attn_q.weight"); + L.wkv_a_mqa = find("attn_kv_a_mqa.weight"); + L.wkv_a_norm = find("attn_kv_a_norm.weight"); + L.wk_b = find("attn_k_b.weight"); + L.wv_b = find("attn_v_b.weight"); + L.wkv_b = find("attn_kv_b.weight"); + L.wqkv_gate = find("attn_gate.weight"); + const bool q_ok = L.wq || (L.wq_a && L.wq_a_norm && L.wq_b); + if (!q_ok || !L.wkv_a_mqa || !L.wkv_a_norm || !L.wqkv_gate || + ((!L.wk_b || !L.wv_b) && !L.wkv_b)) { + return fail("MLA layer " + std::to_string(il) + " is incomplete"); + } + // The first native path intentionally requires absorbed MLA. It is + // the official K3 layout and stores one compact K-only cache. + if (!L.wk_b || !L.wv_b) { + return fail("MLA layer " + std::to_string(il) + + " uses unabsorbed attn_kv_b; not supported by the native cache yet"); + } + } + + if (il < out.n_dense_lead) { + L.ffn_gate = find("ffn_gate.weight"); + L.ffn_up = find("ffn_up.weight"); + L.ffn_down = find("ffn_down.weight"); + if (!L.ffn_gate || !L.ffn_up || !L.ffn_down) { + return fail("dense FFN layer " + std::to_string(il) + " is incomplete"); + } + } else { + L.ffn_gate_inp = find("ffn_gate_inp.weight"); + L.ffn_exp_probs_b = find("exp_probs_b.bias"); + L.ffn_gate_exps = find("ffn_gate_exps.weight"); + L.ffn_up_exps = find("ffn_up_exps.weight"); + L.ffn_down_exps = find("ffn_down_exps.weight"); + L.ffn_routed_down = find("ffn_routed_down.weight"); + L.ffn_routed_up = find("ffn_routed_up.weight"); + L.ffn_routed_norm = find("ffn_routed_norm.weight"); + L.ffn_gate_shexp = find("ffn_gate_shexp.weight"); + L.ffn_up_shexp = find("ffn_up_shexp.weight"); + L.ffn_down_shexp = find("ffn_down_shexp.weight"); + if (!L.ffn_gate_inp || !L.ffn_exp_probs_b || !L.ffn_gate_exps || + !L.ffn_up_exps || !L.ffn_down_exps || !L.ffn_routed_down || + !L.ffn_routed_up || !L.ffn_gate_shexp || !L.ffn_up_shexp || + !L.ffn_down_shexp) { + return fail("latent MoE layer " + std::to_string(il) + " is incomplete"); + } + } + } + + out.buf = ggml_backend_alloc_ctx_tensors(meta_ctx, backend); + if (!out.buf) return fail("unable to allocate resident tensor buffer"); + + GgufMmap mmap; + std::string mmap_error; + if (!mmap.open(path, mmap_error)) return fail(mmap_error); + const auto * base = static_cast(mmap.data()); + const size_t file_size = mmap.size(); + const size_t data_start = gguf_get_data_offset(gctx); + size_t copied = 0; + for (int64_t tid = 0; tid < gguf_get_n_tensors(gctx); ++tid) { + const char * tensor_name = gguf_get_tensor_name(gctx, tid); + ggml_tensor * tensor = ggml_get_tensor(meta_ctx, tensor_name); + if (!tensor) continue; + const size_t offset = gguf_get_tensor_offset(gctx, tid); + const size_t bytes = gguf_get_tensor_size(gctx, tid); + if (!gguf_tensor_in_file(data_start, offset, bytes, file_size)) { + return fail(gguf_bounds_error("Kimi-K3 GGUF", tensor_name, + ggml_type_name(gguf_get_tensor_type(gctx, tid)), data_start, + offset, bytes, file_size)); + } + ggml_backend_tensor_set(tensor, base + data_start + offset, 0, bytes); + copied += bytes; + } + + gguf_free(gctx); + std::fprintf(stderr, + "[kimi-k3] loaded %.2f GiB: layers=%d (KDA=%zu MLA=%zu) hidden=%d " + "experts=%d top=%d latent=%d vocab=%d\n", + static_cast(copied) / (1024.0 * 1024.0 * 1024.0), + out.n_layer, + static_cast(std::count_if(out.layers.begin(), out.layers.end(), + [](const KimiK3Layer & l) { return l.recurrent; })), + static_cast(std::count_if(out.layers.begin(), out.layers.end(), + [](const KimiK3Layer & l) { return !l.recurrent; })), + out.n_embd, out.n_expert, out.n_expert_used, + out.n_expert_latent, out.n_vocab); + std::fflush(stderr); + return true; +} + +void free_kimi_k3_weights(KimiK3Weights & w) { + if (w.buf) ggml_backend_buffer_free(w.buf); + if (w.ctx) ggml_free(w.ctx); + w = KimiK3Weights{}; +} + +} // namespace dflash::common diff --git a/server/test/smoke_kimi_k3_forward.cpp b/server/test/smoke_kimi_k3_forward.cpp new file mode 100644 index 000000000..9c6b42c8f --- /dev/null +++ b/server/test/smoke_kimi_k3_forward.cpp @@ -0,0 +1,70 @@ +#include "kimi_k3/kimi_k3_backend.h" +#include "server/tokenizer.h" + +#include +#include +#include + +using namespace dflash::common; + +int main(int argc, char ** argv) { + if (argc < 2) { + std::fprintf(stderr, + "usage: %s [gpu=0] [n_gen=16] [prompt]\n", + argv[0]); + return 2; + } + const char * model = argv[1]; + const int gpu = argc > 2 ? std::atoi(argv[2]) : 0; + const int n_gen = argc > 3 ? std::atoi(argv[3]) : 16; + const std::string prompt = argc > 4 + ? argv[4] + : "According to all known laws"; + + Tokenizer tokenizer; + if (!tokenizer.load_from_gguf(model)) { + std::fprintf(stderr, "[kimi-k3-smoke] tokenizer load failed\n"); + return 1; + } + std::vector prompt_ids = tokenizer.encode(prompt); + if (prompt_ids.empty()) { + std::fprintf(stderr, "[kimi-k3-smoke] prompt tokenized to zero IDs\n"); + return 1; + } + + KimiK3BackendConfig config; + config.model_path = model; + config.device.gpu = gpu; + config.device.max_ctx = 4096; + KimiK3Backend backend(config); + if (!backend.init()) return 1; + + GenerateRequest request; + request.prompt = prompt_ids; + request.n_gen = n_gen; + request.do_sample = false; + DaemonIO io; + GenerateResult result = backend.generate(request, io); + if (!result.ok()) { + std::fprintf(stderr, "[kimi-k3-smoke] generation failed: %s (%s)\n", + std::string(result.error_code()).c_str(), + std::string(result.error_detail()).c_str()); + return 1; + } + + std::printf("[kimi-k3-smoke] prompt_ids:"); + for (int32_t id : prompt_ids) std::printf(" %d", id); + std::printf("\n[kimi-k3-smoke] output_ids:"); + for (int32_t id : result.tokens) std::printf(" %d", id); + std::printf("\n[kimi-k3-smoke] text: %s%s\n", + prompt.c_str(), tokenizer.decode(result.tokens).c_str()); + const double prefill_rate = result.prefill_s > 0.0 + ? static_cast(prompt_ids.size()) / result.prefill_s : 0.0; + const double decode_rate = result.decode_s > 0.0 + ? static_cast(result.tokens.size()) / result.decode_s : 0.0; + std::printf("[kimi-k3-smoke] prefill=%.3fs (%.2f tok/s) " + "decode=%.3fs (%.2f tok/s)\n", + result.prefill_s, prefill_rate, + result.decode_s, decode_rate); + return 0; +} diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp index b75eee35e..67daef4fe 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -265,9 +265,9 @@ static void test_feature_gate_layer_split_requires_supported_arch() { for (const char * arch : {"qwen35", "laguna", "gemma4", "deepseek4"}) { TEST_ASSERT(gate_result(args, arch, PlacementBackend::Cuda).empty()); } - // These two do not: the factory would hand the split placement to a + // These do not: the factory would hand the split placement to a // monolithic backend, which reads only the primary GPU. - for (const char * arch : {"qwen35moe", "qwen3"}) { + for (const char * arch : {"qwen35moe", "qwen3", "kimi-k3"}) { TEST_ASSERT(!gate_result(args, arch, PlacementBackend::Cuda).empty()); } @@ -276,6 +276,7 @@ static void test_feature_gate_layer_split_requires_supported_arch() { single.model_path = "/nonexistent/model.gguf"; TEST_ASSERT(gate_result(single, "qwen35moe", PlacementBackend::Cuda).empty()); TEST_ASSERT(gate_result(single, "qwen3", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_result(single, "kimi-k3", PlacementBackend::Cuda).empty()); } // ── Inert-flag warnings ───────────────────────────────────────────────── @@ -316,9 +317,10 @@ static void test_feature_warnings_report_inert_draft() { args.model_path = "/nonexistent/model.gguf"; args.draft_path = "/nonexistent/draft.gguf"; - // qwen3 and deepseek4 never forward a draft model. + // These native AR backends never forward a draft model. TEST_ASSERT(warns_about(warn_result(args, "qwen3"), "--draft")); TEST_ASSERT(warns_about(warn_result(args, "deepseek4"), "--draft")); + TEST_ASSERT(warns_about(warn_result(args, "kimi-k3"), "--draft")); // laguna and gemma4 forward it only when monolithic. TEST_ASSERT(!warns_about(warn_result(args, "laguna"), "--draft")); TEST_ASSERT(!warns_about(warn_result(args, "gemma4"), "--draft")); @@ -380,7 +382,7 @@ static void test_model_capability_tables() { // arch_is_supported() must match create_backend()'s dispatch chain. for (const char * arch : {"qwen35", "qwen35moe", "laguna", - "qwen3", "gemma4", "deepseek4"}) { + "qwen3", "gemma4", "deepseek4", "kimi-k3"}) { TEST_ASSERT(arch_is_supported(arch)); } TEST_ASSERT(!arch_is_supported("")); @@ -392,6 +394,7 @@ static void test_model_capability_tables() { TEST_ASSERT(!arch_has_expert_offload("qwen35")); // deepseek4 is mixture-of-experts but has no hot/cold offload path. TEST_ASSERT(!arch_has_expert_offload("deepseek4")); + TEST_ASSERT(!arch_has_expert_offload("kimi-k3")); // Every capability predicate must be false for an architecture the // factory cannot build, so no rule can admit an unbuildable model. From 5c51598bb01c2c3e8373655c83f1cd9a68620f44 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:20:37 +0200 Subject: [PATCH 05/20] fix(moe): pad streamed quant tensors on device --- server/src/common/moe_hybrid_stream.cpp | 386 ++++++++++++++++++++---- server/test/bench_kimi_k3_hetero.cpp | 293 ++++++++++++++---- 2 files changed, 576 insertions(+), 103 deletions(-) diff --git a/server/src/common/moe_hybrid_stream.cpp b/server/src/common/moe_hybrid_stream.cpp index 7f9c42697..ec3b4cc3a 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -291,6 +291,9 @@ class PersistentStreamExpertGraph { const void * gate_data, const void * up_data, const void * down_data, + size_t gate_alloc_bytes, + size_t up_alloc_bytes, + size_t down_alloc_bytes, std::string * err) { destroy(); if (!backend || !expert_buffer || batch <= 0 || !gate_data || !down_data || @@ -336,22 +339,19 @@ class PersistentStreamExpertGraph { auto bind_external = [&](ggml_tensor * tensor, const void * data, + size_t available_bytes, const char * label) -> bool { if (!tensor || !data) { if (err) *err = std::string("invalid streamed ") + label + " tensor binding"; return false; } - // Some GPU quant kernels require row-tail padding beyond - // ggml_nbytes(). The ordinary allocator supplies that padding, - // but a compact streamed record does not. Refuse such an adapter - // until it provides a padded device layout; otherwise a kernel - // could zero or read into the following component. - if (ggml_backend_buffer_get_alloc_size(expert_buffer, tensor) != - ggml_nbytes(tensor)) { + const size_t required_bytes = + ggml_backend_buffer_get_alloc_size(expert_buffer, tensor); + if (available_bytes < required_bytes) { if (err) *err = std::string("streamed ") + label + - " tensor requires backend row padding; use a padded " - "device layout"; + " device allocation is smaller than the backend's " + "padded tensor requirement"; return false; } const size_t alignment = @@ -372,13 +372,14 @@ class PersistentStreamExpertGraph { return true; }; if (spec.fused_gate_up) { - if (!bind_external(gate_up_, gate_data, "gate_up") || - !bind_external(down_, down_data, "down")) { + if (!bind_external(gate_up_, gate_data, gate_alloc_bytes, + "gate_up") || + !bind_external(down_, down_data, down_alloc_bytes, "down")) { return false; } - } else if (!bind_external(gate_, gate_data, "gate") || - !bind_external(up_, up_data, "up") || - !bind_external(down_, down_data, "down")) { + } else if (!bind_external(gate_, gate_data, gate_alloc_bytes, "gate") || + !bind_external(up_, up_data, up_alloc_bytes, "up") || + !bind_external(down_, down_data, down_alloc_bytes, "down")) { return false; } @@ -524,6 +525,29 @@ class PersistentStreamExpertGraph { } // namespace struct MoeHybridStreamEngine::Runtime { + struct DeviceComponentLayout { + MoeExpertComponentKind kind = MoeExpertComponentKind::Gate; + size_t offset = 0; + size_t logical_bytes = 0; + size_t alloc_bytes = 0; + }; + + struct DeviceExpertLayout { + MoeStreamExpertSpec spec{}; + DeviceComponentLayout components[3]{}; + int component_count = 0; + size_t bytes = 0; + bool configured = false; + + const DeviceComponentLayout * component( + MoeExpertComponentKind kind) const { + for (int i = 0; i < component_count; ++i) { + if (components[i].kind == kind) return &components[i]; + } + return nullptr; + } + }; + struct DeviceSlot { void * data = nullptr; cudaEvent_t ready = nullptr; @@ -536,6 +560,7 @@ struct MoeHybridStreamEngine::Runtime { uint64_t last_touch = 0; MoeNvmeLease host_lease; MoeExpertIoLayout layout{}; + DeviceExpertLayout device_layout{}; }; ggml_backend_t backend = nullptr; @@ -550,6 +575,7 @@ struct MoeHybridStreamEngine::Runtime { size_t device_pool_bytes = 0; std::vector device_slots; std::unordered_map device_index; + std::unordered_map layer_device_layouts; uint64_t device_clock = 0; uint64_t device_cache_hits = 0; uint64_t device_cache_misses = 0; @@ -562,8 +588,37 @@ struct MoeHybridStreamEngine::Runtime { }; template -bool allocate_device_cache(RuntimeT & runtime, std::string * err) { - runtime.device_stride = align_up(runtime.max_expert_bytes, 256); +void release_device_cache(RuntimeT & runtime) { + if (runtime.backend) ggml_backend_synchronize(runtime.backend); + if (runtime.transfer_stream) { + (void) cudaStreamSynchronize(runtime.transfer_stream); + } + runtime.graph_cache.clear(); + for (auto & slot : runtime.device_slots) { + slot.host_lease.reset(); + if (slot.ready) (void) cudaEventDestroy(slot.ready); + slot.ready = nullptr; + slot.data = nullptr; + slot.pending = false; + } + runtime.device_slots.clear(); + runtime.device_index.clear(); + runtime.active_slot = -1; + if (runtime.device_pool_buffer) { + ggml_backend_buffer_free(runtime.device_pool_buffer); + } + runtime.device_pool_buffer = nullptr; + runtime.device_pool = nullptr; + runtime.device_stride = 0; + runtime.device_pool_bytes = 0; +} + +template +bool allocate_device_cache(RuntimeT & runtime, std::string * err, + size_t minimum_stride = 0) { + const size_t logical_stride = + std::max(runtime.max_expert_bytes, minimum_stride); + runtime.device_stride = align_up(logical_stride, 256); if (runtime.device_stride == 0) { if (err) *err = "SSD device-cache stride overflow"; return false; @@ -619,7 +674,130 @@ bool allocate_device_cache(RuntimeT & runtime, std::string * err) { for (size_t i = 0; i < runtime.device_slots.size(); ++i) { runtime.device_slots[i].data = base + i * runtime.device_stride; } - runtime.config.device_slots = (int) attempt_slots; + return true; +} + +template +bool build_device_expert_layout( + RuntimeT & runtime, + const MoeStreamExpertSpec & spec, + typename RuntimeT::DeviceExpertLayout & out, + std::string * err) { + out = typename RuntimeT::DeviceExpertLayout{}; + if (spec.input_dim <= 0 || spec.intermediate_dim <= 0 || + spec.output_dim <= 0 || !valid_ggml_type(spec.down_type) || + (spec.fused_gate_up + ? !valid_ggml_type(spec.gate_up_type) + : (!valid_ggml_type(spec.gate_type) || + !valid_ggml_type(spec.up_type)))) { + if (err) *err = "invalid streamed expert shape or tensor type"; + return false; + } + + ggml_backend_buffer_type_t buft = + ggml_backend_get_default_buffer_type(runtime.backend); + const size_t alignment = ggml_backend_buft_get_alignment(buft); + if (alignment == 0 || (alignment & (alignment - 1)) != 0) { + if (err) *err = "streamed expert backend alignment is invalid"; + return false; + } + + ggml_init_params params{}; + params.mem_size = 128 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + if (err) *err = "ggml_init failed for streamed device layout"; + return false; + } + + size_t cursor = 0; + auto add = [&](MoeExpertComponentKind kind, ggml_type type, + int64_t columns, int64_t rows, const char * label) -> bool { + if (out.component_count >= 3 || columns <= 0 || rows <= 0 || + columns % ggml_blck_size(type) != 0) { + if (err) *err = std::string("invalid streamed ") + label + + " tensor dimensions"; + return false; + } + ggml_tensor * tensor = + ggml_new_tensor_2d(ctx, type, columns, rows); + const size_t logical_bytes = ggml_nbytes(tensor); + const size_t alloc_bytes = + ggml_backend_buft_get_alloc_size(buft, tensor); + const size_t offset = align_up(cursor, alignment); + if (offset == 0 && cursor != 0) { + if (err) *err = "streamed device component alignment overflow"; + return false; + } + if (alloc_bytes < logical_bytes || + offset > std::numeric_limits::max() - alloc_bytes) { + if (err) *err = "streamed device component size overflow"; + return false; + } + out.components[out.component_count++] = { + kind, offset, logical_bytes, alloc_bytes}; + cursor = offset + alloc_bytes; + return true; + }; + + bool ok = true; + if (spec.fused_gate_up) { + ok = spec.intermediate_dim <= std::numeric_limits::max() / 2 && + add(MoeExpertComponentKind::FusedGateUp, spec.gate_up_type, + spec.input_dim, 2LL * spec.intermediate_dim, "gate_up") && + add(MoeExpertComponentKind::Down, spec.down_type, + spec.intermediate_dim, spec.output_dim, "down"); + } else { + ok = add(MoeExpertComponentKind::Gate, spec.gate_type, + spec.input_dim, spec.intermediate_dim, "gate") && + add(MoeExpertComponentKind::Up, spec.up_type, + spec.input_dim, spec.intermediate_dim, "up") && + add(MoeExpertComponentKind::Down, spec.down_type, + spec.intermediate_dim, spec.output_dim, "down"); + } + if (ok) { + out.bytes = align_up(cursor, std::max(256, alignment)); + if (out.bytes == 0) { + if (err) *err = "streamed device expert stride overflow"; + ok = false; + } + } + ggml_free(ctx); + if (!ok) return false; + out.spec = spec; + out.configured = true; + return true; +} + +template +bool prepare_device_expert_layout(RuntimeT & runtime, int layer, + const MoeStreamExpertSpec & spec, + std::string * err) { + auto existing = runtime.layer_device_layouts.find(layer); + if (existing != runtime.layer_device_layouts.end()) { + if (!same_stream_spec(existing->second.spec, spec)) { + if (err) *err = "streamed expert specification changed within one layer"; + return false; + } + return true; + } + + typename RuntimeT::DeviceExpertLayout layout; + if (!build_device_expert_layout(runtime, spec, layout, err)) return false; + const bool has_compact_cached_slot = std::any_of( + runtime.device_slots.begin(), runtime.device_slots.end(), + [&](const auto & slot) { + return slot.valid && slot.key.layer == layer && + !slot.device_layout.configured; + }); + if (layout.bytes > runtime.device_stride || has_compact_cached_slot) { + const size_t required_stride = + std::max(layout.bytes, runtime.device_stride); + release_device_cache(runtime); + if (!allocate_device_cache(runtime, err, required_stride)) return false; + } + runtime.layer_device_layouts.emplace(layer, layout); return true; } @@ -932,6 +1110,8 @@ bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, dst.valid = false; dst.cache_managed = false; dst.key = {}; + dst.layout = MoeExpertIoLayout{}; + dst.device_layout = Runtime::DeviceExpertLayout{}; if (!dst.ready) { const cudaError_t event_create_err = @@ -949,18 +1129,92 @@ bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, if (err) *err = "streamed expert exceeds GPU device slot"; return false; } - for (int i = 0; i < lease.layout().span_count; ++i) { - const MoeExpertIoSpan & span = lease.layout().spans[i]; - cudaError_t gpu_err = cudaMemcpyAsync( - static_cast(dst.data) + span.device_offset, - lease.data() + span.buffer_offset, - span.bytes, cudaMemcpyHostToDevice, runtime_->transfer_stream); - if (gpu_err != cudaSuccess) { - (void) cudaStreamSynchronize(runtime_->transfer_stream); - if (err) *err = std::string("asynchronous expert H2D failed: ") + - cudaGetErrorString(gpu_err); + const auto prepared = runtime_->layer_device_layouts.find(layer); + if (prepared != runtime_->layer_device_layouts.end()) { + const Runtime::DeviceExpertLayout & device_layout = prepared->second; + if (!device_layout.configured || + device_layout.bytes > runtime_->device_stride) { + if (err) *err = "prepared streamed expert exceeds GPU device stride"; return false; } + for (int i = 0; i < device_layout.component_count; ++i) { + const Runtime::DeviceComponentLayout & device_component = + device_layout.components[i]; + const MoeExpertComponentLayout * io_component = + lease.layout().component(device_component.kind); + if (!io_component || + io_component->bytes != device_component.logical_bytes) { + if (err) *err = "streamed device and storage components disagree"; + return false; + } + + const uint8_t * source = nullptr; + for (int span_index = 0; + span_index < lease.layout().span_count; ++span_index) { + const MoeExpertIoSpan & span = + lease.layout().spans[span_index]; + if (io_component->device_offset < span.device_offset) continue; + const size_t delta = + io_component->device_offset - span.device_offset; + if (delta <= span.bytes && + io_component->bytes <= span.bytes - delta) { + source = lease.data() + span.buffer_offset + delta; + break; + } + } + if (!source) { + if (err) *err = "streamed component is not contained in its I/O span"; + return false; + } + + cudaError_t gpu_err = cudaMemcpyAsync( + static_cast(dst.data) + device_component.offset, + source, device_component.logical_bytes, + cudaMemcpyHostToDevice, runtime_->transfer_stream); + if (gpu_err == cudaSuccess && + device_component.alloc_bytes > device_component.logical_bytes) { + gpu_err = cudaMemsetAsync( + static_cast(dst.data) + device_component.offset + + device_component.logical_bytes, + 0, + device_component.alloc_bytes - + device_component.logical_bytes, + runtime_->transfer_stream); + } + if (gpu_err != cudaSuccess) { + (void) cudaStreamSynchronize(runtime_->transfer_stream); + if (err) *err = std::string("asynchronous expert H2D failed: ") + + cudaGetErrorString(gpu_err); + return false; + } + } + dst.device_layout = device_layout; + } else { + // Storage-only callers do not supply a compute specification. Preserve + // their compact byte-for-byte staging contract; numerical evaluation + // always registers an exact backend-padded layout before reaching here. + for (int i = 0; i < lease.layout().span_count; ++i) { + const MoeExpertIoSpan & span = lease.layout().spans[i]; + cudaError_t gpu_err = cudaMemcpyAsync( + static_cast(dst.data) + span.device_offset, + lease.data() + span.buffer_offset, + span.bytes, cudaMemcpyHostToDevice, runtime_->transfer_stream); + if (gpu_err != cudaSuccess) { + (void) cudaStreamSynchronize(runtime_->transfer_stream); + if (err) *err = std::string("asynchronous expert H2D failed: ") + + cudaGetErrorString(gpu_err); + return false; + } + } + dst.device_layout.component_count = lease.layout().component_count; + dst.device_layout.bytes = lease.layout().payload_bytes; + for (int i = 0; i < lease.layout().component_count; ++i) { + const MoeExpertComponentLayout & component = + lease.layout().components[i]; + dst.device_layout.components[i] = { + component.kind, component.device_offset, + component.bytes, component.bytes}; + } } const cudaError_t event_err = cudaEventRecord(dst.ready, runtime_->transfer_stream); if (event_err != cudaSuccess) { @@ -1137,9 +1391,10 @@ const void * MoeHybridStreamEngine::scratch_gate_data() const { const Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) runtime_->active_slot]; const MoeExpertComponentKind kind = slot.layout.fused_gate_up ? MoeExpertComponentKind::FusedGateUp : MoeExpertComponentKind::Gate; - const MoeExpertComponentLayout * component = slot.layout.component(kind); + const Runtime::DeviceComponentLayout * component = + slot.device_layout.component(kind); return component - ? static_cast(slot.data) + component->device_offset + ? static_cast(slot.data) + component->offset : nullptr; } @@ -1147,50 +1402,51 @@ const void * MoeHybridStreamEngine::scratch_up_data() const { if (!runtime_ || runtime_->active_slot < 0) return nullptr; const Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) runtime_->active_slot]; if (slot.layout.fused_gate_up) return nullptr; - const MoeExpertComponentLayout * component = - slot.layout.component(MoeExpertComponentKind::Up); + const Runtime::DeviceComponentLayout * component = + slot.device_layout.component(MoeExpertComponentKind::Up); return component - ? static_cast(slot.data) + component->device_offset + ? static_cast(slot.data) + component->offset : nullptr; } const void * MoeHybridStreamEngine::scratch_down_data() const { if (!runtime_ || runtime_->active_slot < 0) return nullptr; const Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) runtime_->active_slot]; - const MoeExpertComponentLayout * component = - slot.layout.component(MoeExpertComponentKind::Down); + const Runtime::DeviceComponentLayout * component = + slot.device_layout.component(MoeExpertComponentKind::Down); return component - ? static_cast(slot.data) + component->device_offset + ? static_cast(slot.data) + component->offset : nullptr; } size_t MoeHybridStreamEngine::scratch_gate_bytes() const { if (!runtime_ || runtime_->active_slot < 0) return 0; - const MoeExpertIoLayout & layout = - runtime_->device_slots[(size_t) runtime_->active_slot].layout; - const MoeExpertComponentKind kind = layout.fused_gate_up + const Runtime::DeviceSlot & slot = + runtime_->device_slots[(size_t) runtime_->active_slot]; + const MoeExpertComponentKind kind = slot.layout.fused_gate_up ? MoeExpertComponentKind::FusedGateUp : MoeExpertComponentKind::Gate; - const MoeExpertComponentLayout * component = layout.component(kind); - return component ? component->bytes : 0; + const Runtime::DeviceComponentLayout * component = + slot.device_layout.component(kind); + return component ? component->logical_bytes : 0; } size_t MoeHybridStreamEngine::scratch_up_bytes() const { if (!runtime_ || runtime_->active_slot < 0) return 0; - const MoeExpertIoLayout & layout = - runtime_->device_slots[(size_t) runtime_->active_slot].layout; - if (layout.fused_gate_up) return 0; - const MoeExpertComponentLayout * component = - layout.component(MoeExpertComponentKind::Up); - return component ? component->bytes : 0; + const Runtime::DeviceSlot & slot = + runtime_->device_slots[(size_t) runtime_->active_slot]; + if (slot.layout.fused_gate_up) return 0; + const Runtime::DeviceComponentLayout * component = + slot.device_layout.component(MoeExpertComponentKind::Up); + return component ? component->logical_bytes : 0; } size_t MoeHybridStreamEngine::scratch_down_bytes() const { if (!runtime_ || runtime_->active_slot < 0) return 0; - const MoeExpertIoLayout & layout = - runtime_->device_slots[(size_t) runtime_->active_slot].layout; - const MoeExpertComponentLayout * component = - layout.component(MoeExpertComponentKind::Down); - return component ? component->bytes : 0; + const Runtime::DeviceSlot & slot = + runtime_->device_slots[(size_t) runtime_->active_slot]; + const Runtime::DeviceComponentLayout * component = + slot.device_layout.component(MoeExpertComponentKind::Down); + return component ? component->logical_bytes : 0; } size_t MoeHybridStreamEngine::pinned_bytes() const { @@ -1509,6 +1765,10 @@ bool eval_moe_streamed_experts( auto & runtime = *engine.runtime_; std::lock_guard compute_guard(runtime.compute_mutex); + if (!prepare_device_expert_layout( + runtime, batch.layer, spec, err)) { + return false; + } size_t route_slots = 0; if (!checked_mul_size((size_t) batch.top_k, @@ -1575,6 +1835,24 @@ bool eval_moe_streamed_experts( const MoeExpertIoLayout & layout = runtime.device_slots[(size_t) active].layout; if (!validate_moe_stream_expert_layout(spec, layout, err)) return false; + const auto & device_layout = + runtime.device_slots[(size_t) active].device_layout; + const MoeExpertComponentKind gate_kind = spec.fused_gate_up + ? MoeExpertComponentKind::FusedGateUp + : MoeExpertComponentKind::Gate; + const auto * gate_component = + device_layout.component(gate_kind); + const auto * up_component = + spec.fused_gate_up + ? nullptr + : device_layout.component(MoeExpertComponentKind::Up); + const auto * down_component = + device_layout.component(MoeExpertComponentKind::Down); + if (!gate_component || !down_component || + (!spec.fused_gate_up && !up_component)) { + if (err) *err = "streamed device layout is missing an expert component"; + return false; + } std::unique_ptr built( new (std::nothrow) PersistentStreamExpertGraph); @@ -1586,7 +1864,11 @@ bool eval_moe_streamed_experts( spec, graph_batch, engine.scratch_gate_data(), engine.scratch_up_data(), - engine.scratch_down_data(), err)) { + engine.scratch_down_data(), + gate_component->alloc_bytes, + up_component ? up_component->alloc_bytes : 0, + down_component->alloc_bytes, + err)) { return false; } built->last_touch = touch; diff --git a/server/test/bench_kimi_k3_hetero.cpp b/server/test/bench_kimi_k3_hetero.cpp index 581498a41..62dabb0d4 100644 --- a/server/test/bench_kimi_k3_hetero.cpp +++ b/server/test/bench_kimi_k3_hetero.cpp @@ -1,15 +1,17 @@ // Read-only Kimi-K3 routed-core qualification for heterogeneous Lucebox. // -// This is not a substitute for the Kimi model graph. It isolates the part -// whose placement is genuinely new: exact routed experts moving from NVMe to -// one GPU while that GPU evaluates the persistent IQ1_S + SiTU expert graph. -// The source file only supplies bytes and is never modified. +// It accepts either a real Kimi-K3 GGUF (preferred) or a raw source file large +// enough to emulate the released 2.8T geometry. With a GGUF, tensor types, +// dimensions, layer offsets, and per-expert strides come from model metadata; +// the bytes evaluated by the common stream engine are the actual checkpoint +// weights. The source file is read-only and is never modified. #include "common/moe_hybrid_stream.h" #include "ggml-backend.h" #include "ggml-cuda.h" #include "ggml.h" +#include "gguf.h" #include #include @@ -61,15 +63,16 @@ double gib(uint64_t bytes) { return (double) bytes / (1024.0 * 1024.0 * 1024.0); } -std::vector route_for(int token, int layer, int top_k, bool repeat) { - std::vector population((size_t) kKimiExperts); +std::vector route_for(int token, int layer, int n_expert, + int top_k, bool repeat) { + std::vector population((size_t) n_expert); std::iota(population.begin(), population.end(), 0); const uint64_t token_key = repeat ? 0 : (uint64_t) token; std::mt19937_64 rng( 0x4b494d49334c5543ULL ^ (token_key * 0x9e3779b97f4a7c15ULL) ^ ((uint64_t) layer * 0xbf58476d1ce4e5b9ULL)); for (int i = 0; i < top_k; ++i) { - std::uniform_int_distribution choose(i, kKimiExperts - 1); + std::uniform_int_distribution choose(i, n_expert - 1); const int selected = choose(rng); std::swap(population[(size_t) i], population[(size_t) selected]); } @@ -77,21 +80,171 @@ std::vector route_for(int token, int layer, int top_k, bool repeat) { return population; } +struct KimiGgufStreamLayout { + bool detected = false; + int n_expert = kKimiExperts; + int top_k = kKimiTopK; + int latent = (int) kKimiLatent; + int expert_ff = (int) kKimiExpertFf; + float situ_beta = kSituBeta; + float situ_linear_beta = kSituLinearBeta; + ggml_type gate_type = GGML_TYPE_IQ1_S; + ggml_type up_type = GGML_TYPE_IQ1_S; + ggml_type down_type = GGML_TYPE_IQ1_S; + size_t expert_bytes = 0; + std::vector regions; +}; + +uint32_t gguf_u32_or(const gguf_context * g, const char * key, uint32_t fallback) { + const int64_t id = gguf_find_key(g, key); + return id >= 0 && gguf_get_kv_type(g, id) == GGUF_TYPE_UINT32 + ? gguf_get_val_u32(g, id) : fallback; +} + +float gguf_f32_or(const gguf_context * g, const char * key, float fallback) { + const int64_t id = gguf_find_key(g, key); + return id >= 0 && gguf_get_kv_type(g, id) == GGUF_TYPE_FLOAT32 + ? gguf_get_val_f32(g, id) : fallback; +} + +bool inspect_kimi_gguf_layout(const char * path, + KimiGgufStreamLayout & out, + std::string & error) { + ggml_context * meta = nullptr; + gguf_init_params params{}; + params.no_alloc = true; + params.ctx = &meta; + gguf_context * g = gguf_init_from_file(path, params); + if (!g || !meta) { + if (g) gguf_free(g); + if (meta) ggml_free(meta); + return false; // Non-GGUF inputs retain the synthetic compatibility path. + } + const int64_t arch_id = gguf_find_key(g, "general.architecture"); + const char * arch = arch_id >= 0 ? gguf_get_val_str(g, arch_id) : ""; + if (std::strcmp(arch, "kimi-k3") != 0) { + gguf_free(g); + ggml_free(meta); + return false; + } + + out.detected = true; + out.n_expert = (int) gguf_u32_or(g, "kimi-k3.expert_count", 0); + out.top_k = (int) gguf_u32_or(g, "kimi-k3.expert_used_count", 0); + out.latent = (int) gguf_u32_or(g, "kimi-k3.expert_latent_length", 0); + out.expert_ff = (int) gguf_u32_or(g, "kimi-k3.expert_feed_forward_length", 0); + out.situ_beta = gguf_f32_or(g, "kimi-k3.activation.situ_beta", kSituBeta); + out.situ_linear_beta = gguf_f32_or( + g, "kimi-k3.activation.situ_linear_beta", kSituLinearBeta); + const int n_layer = (int) gguf_u32_or(g, "kimi-k3.block_count", 0); + const int dense_lead = (int) gguf_u32_or( + g, "kimi-k3.leading_dense_block_count", 0); + if (out.n_expert <= 0 || out.top_k <= 0 || out.top_k > out.n_expert || + out.latent <= 0 || out.expert_ff <= 0 || n_layer <= dense_lead) { + error = "invalid Kimi-K3 GGUF routing metadata"; + gguf_free(g); + ggml_free(meta); + return false; + } + + const size_t data_start = gguf_get_data_offset(g); + for (int il = dense_lead; il < n_layer; ++il) { + char gate_name[128], up_name[128], down_name[128]; + std::snprintf(gate_name, sizeof(gate_name), + "blk.%d.ffn_gate_exps.weight", il); + std::snprintf(up_name, sizeof(up_name), + "blk.%d.ffn_up_exps.weight", il); + std::snprintf(down_name, sizeof(down_name), + "blk.%d.ffn_down_exps.weight", il); + const int64_t gate_id = gguf_find_tensor(g, gate_name); + const int64_t up_id = gguf_find_tensor(g, up_name); + const int64_t down_id = gguf_find_tensor(g, down_name); + ggml_tensor * gate = ggml_get_tensor(meta, gate_name); + ggml_tensor * up = ggml_get_tensor(meta, up_name); + ggml_tensor * down = ggml_get_tensor(meta, down_name); + if (gate_id < 0 || up_id < 0 || down_id < 0 || !gate || !up || !down || + gate->ne[0] != out.latent || gate->ne[1] != out.expert_ff || + gate->ne[2] != out.n_expert || + up->ne[0] != out.latent || up->ne[1] != out.expert_ff || + up->ne[2] != out.n_expert || + down->ne[0] != out.expert_ff || down->ne[1] != out.latent || + down->ne[2] != out.n_expert) { + error = "Kimi-K3 expert tensor shape mismatch at model layer " + + std::to_string(il); + gguf_free(g); + ggml_free(meta); + return false; + } + + const size_t gate_size = gguf_get_tensor_size(g, gate_id); + const size_t up_size = gguf_get_tensor_size(g, up_id); + const size_t down_size = gguf_get_tensor_size(g, down_id); + if (gate_size % (size_t) out.n_expert != 0 || + up_size % (size_t) out.n_expert != 0 || + down_size % (size_t) out.n_expert != 0) { + error = "Kimi-K3 expert tensor is not expert-major stridable at layer " + + std::to_string(il); + gguf_free(g); + ggml_free(meta); + return false; + } + const ggml_type gt = gguf_get_tensor_type(g, gate_id); + const ggml_type ut = gguf_get_tensor_type(g, up_id); + const ggml_type dt = gguf_get_tensor_type(g, down_id); + if (out.regions.empty()) { + out.gate_type = gt; + out.up_type = ut; + out.down_type = dt; + } else if (out.gate_type != gt || out.up_type != ut || out.down_type != dt) { + error = "mixed expert tensor types need per-layer stream specs"; + gguf_free(g); + ggml_free(meta); + return false; + } + + LayerExpertRegions regions; + regions.fused_gate_up = false; + regions.expert_bytes_gate = gate_size / (size_t) out.n_expert; + regions.expert_bytes_up = up_size / (size_t) out.n_expert; + regions.expert_bytes_down = down_size / (size_t) out.n_expert; + regions.gate_exps = { + data_start + gguf_get_tensor_offset(g, gate_id), gate_size, 0}; + regions.up_exps = { + data_start + gguf_get_tensor_offset(g, up_id), up_size, 0}; + regions.down_exps = { + data_start + gguf_get_tensor_offset(g, down_id), down_size, 0}; + const size_t bytes = regions.expert_bytes_gate + + regions.expert_bytes_up + + regions.expert_bytes_down; + if (out.expert_bytes == 0) out.expert_bytes = bytes; + if (out.expert_bytes != bytes) { + error = "variable per-layer expert bytes need per-layer stream slots"; + gguf_free(g); + ggml_free(meta); + return false; + } + out.regions.push_back(regions); + } + gguf_free(g); + ggml_free(meta); + return !out.regions.empty(); +} + } // namespace int main(int argc, char ** argv) { if (argc < 2 || argc > 8) { std::fprintf(stderr, - "usage: %s MODEL_FILE [device=1] [tokens=1] [layers=92] " - "[top_k=16] [compute=1] [repeat_routes=0]\n", + "usage: %s MODEL_FILE [device=1] [tokens=1] [layers=all] " + "[top_k=model] [compute=1] [repeat_routes=0]\n", argv[0]); return 2; } uint64_t device_arg = 1; uint64_t tokens_arg = 1; - uint64_t layers_arg = kKimiMoeLayers; - uint64_t top_k_arg = kKimiTopK; + uint64_t layers_arg = 0; + uint64_t top_k_arg = 0; uint64_t compute_arg = 1; uint64_t repeat_arg = 0; uint64_t * values[] = { @@ -104,9 +257,7 @@ int main(int argc, char ** argv) { return 2; } } - if (tokens_arg == 0 || layers_arg == 0 || layers_arg > kKimiMoeLayers || - top_k_arg == 0 || top_k_arg > kKimiExperts || compute_arg > 1 || - repeat_arg > 1) { + if (tokens_arg == 0 || compute_arg > 1 || repeat_arg > 1) { std::fprintf(stderr, "Kimi benchmark arguments are out of range\n"); return 2; } @@ -129,32 +280,68 @@ int main(int argc, char ** argv) { } const size_t file_bytes = (size_t) st.st_size; - const size_t gate_bytes = - ggml_row_size(GGML_TYPE_IQ1_S, kKimiLatent) * (size_t) kKimiExpertFf; - const size_t up_bytes = gate_bytes; - const size_t down_bytes = - ggml_row_size(GGML_TYPE_IQ1_S, kKimiExpertFf) * (size_t) kKimiLatent; - const size_t expert_bytes = gate_bytes + up_bytes + down_bytes; - const size_t gate_stack = gate_bytes * (size_t) kKimiExperts; - const size_t up_stack = up_bytes * (size_t) kKimiExperts; - const size_t down_stack = down_bytes * (size_t) kKimiExperts; - const size_t required_bytes = gate_stack + up_stack + down_stack; - if (file_bytes < required_bytes) { - std::fprintf(stderr, - "input needs at least %.3f GiB for one Kimi expert stack\n", - gib(required_bytes)); + KimiGgufStreamLayout layout; + std::string layout_error; + const bool actual_kimi = inspect_kimi_gguf_layout(argv[1], layout, layout_error); + if (!actual_kimi && layout.detected) { + std::fprintf(stderr, "Kimi GGUF layout inspection failed: %s\n", + layout_error.c_str()); return 2; } - LayerExpertRegions one_layer; - one_layer.fused_gate_up = false; - one_layer.expert_bytes_gate = gate_bytes; - one_layer.expert_bytes_up = up_bytes; - one_layer.expert_bytes_down = down_bytes; - one_layer.gate_exps = {0, gate_stack}; - one_layer.up_exps = {gate_stack, up_stack}; - one_layer.down_exps = {gate_stack + up_stack, down_stack}; - std::vector regions((size_t) layers_arg, one_layer); + int n_expert = layout.n_expert; + int latent = layout.latent; + int expert_ff = layout.expert_ff; + ggml_type gate_type = layout.gate_type; + ggml_type up_type = layout.up_type; + ggml_type down_type = layout.down_type; + float situ_beta = layout.situ_beta; + float situ_linear_beta = layout.situ_linear_beta; + size_t expert_bytes = layout.expert_bytes; + std::vector regions; + + if (actual_kimi) { + if (layers_arg == 0) layers_arg = layout.regions.size(); + if (top_k_arg == 0) top_k_arg = (uint64_t) layout.top_k; + if (layers_arg > layout.regions.size()) { + std::fprintf(stderr, "requested layers exceed Kimi GGUF MoE layers\n"); + return 2; + } + regions.assign(layout.regions.begin(), + layout.regions.begin() + (size_t) layers_arg); + } else { + if (layers_arg == 0) layers_arg = kKimiMoeLayers; + if (top_k_arg == 0) top_k_arg = kKimiTopK; + const size_t gate_bytes = + ggml_row_size(gate_type, latent) * (size_t) expert_ff; + const size_t up_bytes = gate_bytes; + const size_t down_bytes = + ggml_row_size(down_type, expert_ff) * (size_t) latent; + expert_bytes = gate_bytes + up_bytes + down_bytes; + const size_t gate_stack = gate_bytes * (size_t) n_expert; + const size_t up_stack = up_bytes * (size_t) n_expert; + const size_t down_stack = down_bytes * (size_t) n_expert; + const size_t required_bytes = gate_stack + up_stack + down_stack; + if (file_bytes < required_bytes) { + std::fprintf(stderr, + "input needs at least %.3f GiB for one Kimi expert stack\n", + gib(required_bytes)); + return 2; + } + LayerExpertRegions one_layer; + one_layer.fused_gate_up = false; + one_layer.expert_bytes_gate = gate_bytes; + one_layer.expert_bytes_up = up_bytes; + one_layer.expert_bytes_down = down_bytes; + one_layer.gate_exps = {0, gate_stack}; + one_layer.up_exps = {gate_stack, up_stack}; + one_layer.down_exps = {gate_stack + up_stack, down_stack}; + regions.assign((size_t) layers_arg, one_layer); + } + if (layers_arg == 0 || top_k_arg == 0 || top_k_arg > (uint64_t) n_expert) { + std::fprintf(stderr, "Kimi benchmark routing arguments are out of range\n"); + return 2; + } ggml_backend_t backend = ggml_backend_cuda_init((int) device_arg); if (!backend) { @@ -181,17 +368,17 @@ int main(int argc, char ** argv) { } MoeStreamExpertSpec expert_spec; - expert_spec.input_dim = (int) kKimiLatent; - expert_spec.intermediate_dim = (int) kKimiExpertFf; - expert_spec.output_dim = (int) kKimiLatent; - expert_spec.gate_type = GGML_TYPE_IQ1_S; - expert_spec.up_type = GGML_TYPE_IQ1_S; - expert_spec.down_type = GGML_TYPE_IQ1_S; + expert_spec.input_dim = latent; + expert_spec.intermediate_dim = expert_ff; + expert_spec.output_dim = latent; + expert_spec.gate_type = gate_type; + expert_spec.up_type = up_type; + expert_spec.down_type = down_type; expert_spec.gated_activation = MoeGatedActivation::Situ; - expert_spec.situ_beta = kSituBeta; - expert_spec.situ_linear_beta = kSituLinearBeta; + expert_spec.situ_beta = situ_beta; + expert_spec.situ_linear_beta = situ_linear_beta; - std::vector model_input((size_t) kKimiLatent); + std::vector model_input((size_t) latent); for (size_t i = 0; i < model_input.size(); ++i) { model_input[i] = 0.01f * std::sin((float) i * 0.013f); } @@ -205,11 +392,11 @@ int main(int argc, char ** argv) { for (int token = 0; token < (int) tokens_arg; ++token) { for (int layer = 0; layer < (int) layers_arg; ++layer) { const std::vector experts = route_for( - token, layer, (int) top_k_arg, repeat_arg != 0); + token, layer, n_expert, (int) top_k_arg, repeat_arg != 0); if (compute_arg) { MoeStreamRouteBatch batch; batch.layer = layer; - batch.n_expert = kKimiExperts; + batch.n_expert = n_expert; batch.top_k = (int) top_k_arg; batch.n_tokens = 1; batch.inputs = model_input.data(); @@ -273,10 +460,14 @@ int main(int argc, char ** argv) { "kimi_k3 device=%" PRIu64 " description=%s latent=%lld ff=%lld " "experts=%d top_k=%" PRIu64 " layers=%" PRIu64 " tokens=%" PRIu64 " compute=%s repeat_routes=%s\n", - device_arg, description, (long long) kKimiLatent, - (long long) kKimiExpertFf, kKimiExperts, top_k_arg, layers_arg, - tokens_arg, compute_arg ? "situ-iq1s" : "off", + device_arg, description, (long long) latent, + (long long) expert_ff, n_expert, top_k_arg, layers_arg, + tokens_arg, compute_arg ? "situ" : "off", repeat_arg ? "yes" : "no"); + std::printf("source=%s gate_type=%s up_type=%s down_type=%s\n", + actual_kimi ? "actual-kimi-gguf" : "synthetic-layout", + ggml_type_name(gate_type), ggml_type_name(up_type), + ggml_type_name(down_type)); std::printf( "expert_bytes=%zu expert_mib=%.6f accesses=%" PRIu64 " elapsed_s=%.6f routed_core_tok_s=%.6f experts_s=%.2f\n", From cc1146d4f6acda57b56f214f71c9abc06cff1fd5 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:26:47 +0200 Subject: [PATCH 06/20] test(moe): cover padded MXFP4 stream slots --- server/CMakeLists.txt | 3 +- server/test/test_moe_stream_compute.cpp | 156 +++++++++++++++++++++++- 2 files changed, 154 insertions(+), 5 deletions(-) diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 34237f1cb..eaa0f1659 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -735,7 +735,8 @@ if(DFLASH27B_TESTS) target_include_directories(test_moe_stream_compute PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src ${CMAKE_CURRENT_SOURCE_DIR}/src/common - ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src) target_link_libraries(test_moe_stream_compute PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") diff --git a/server/test/test_moe_stream_compute.cpp b/server/test/test_moe_stream_compute.cpp index b8006e5e0..bbf2ab7cb 100644 --- a/server/test/test_moe_stream_compute.cpp +++ b/server/test/test_moe_stream_compute.cpp @@ -2,6 +2,7 @@ #include "common/moe_hybrid_stream.h" #include "ggml-cuda.h" +#include "ggml-quants.h" #include #include @@ -27,9 +28,12 @@ namespace { struct MoeStreamComputeFixture {}; constexpr int kExperts = 3; -constexpr int kInput = 16; -constexpr int kFf = 24; -constexpr int kOutput = 12; +// 256 deliberately does not satisfy CUDA/HIP's 512-element quantized matrix +// row padding. The MXFP4 case below therefore exercises the padded GPU-slot +// path that real Kimi-K3 exposed. +constexpr int kInput = 256; +constexpr int kFf = 64; +constexpr int kOutput = 128; constexpr int kTokens = 2; constexpr int kTopK = 2; @@ -162,6 +166,73 @@ ModelBytes make_model_bytes(bool expert_major, return model; } +std::vector quantize_mxfp4(const std::vector & values, + int columns, int rows) { + const size_t bytes = ggml_row_size(GGML_TYPE_MXFP4, columns) * + static_cast(rows); + std::vector quantized(bytes); + const size_t written = ggml_quantize_chunk( + GGML_TYPE_MXFP4, values.data(), quantized.data(), 0, + rows, columns, nullptr); + STREAM_REQUIRE(written == bytes); + return quantized; +} + +std::vector dequantize_mxfp4(const std::vector & values, + int columns, int rows) { + const size_t row_bytes = ggml_row_size(GGML_TYPE_MXFP4, columns); + STREAM_REQUIRE(values.size() == row_bytes * static_cast(rows)); + std::vector dequantized( + static_cast(columns) * static_cast(rows)); + for (int row = 0; row < rows; ++row) { + dequantize_row_mxfp4( + reinterpret_cast( + values.data() + static_cast(row) * row_bytes), + dequantized.data() + static_cast(row) * columns, + columns); + } + return dequantized; +} + +ModelBytes make_mxfp4_model_bytes(const std::vector & gate, + const std::vector & up, + const std::vector & down, + std::vector & gate_dequantized, + std::vector & up_dequantized, + std::vector & down_dequantized) { + const std::vector gate_q = quantize_mxfp4( + gate, kInput, kExperts * kFf); + const std::vector up_q = quantize_mxfp4( + up, kInput, kExperts * kFf); + const std::vector down_q = quantize_mxfp4( + down, kFf, kExperts * kOutput); + gate_dequantized = dequantize_mxfp4( + gate_q, kInput, kExperts * kFf); + up_dequantized = dequantize_mxfp4( + up_q, kInput, kExperts * kFf); + down_dequantized = dequantize_mxfp4( + down_q, kFf, kExperts * kOutput); + + ModelBytes model; + model.regions.expert_bytes_gate = + ggml_row_size(GGML_TYPE_MXFP4, kInput) * kFf; + model.regions.expert_bytes_up = model.regions.expert_bytes_gate; + model.regions.expert_bytes_down = + ggml_row_size(GGML_TYPE_MXFP4, kFf) * kOutput; + model.regions.gate_exps = {0, gate_q.size()}; + model.regions.up_exps = {gate_q.size(), up_q.size()}; + model.regions.down_exps = { + gate_q.size() + up_q.size(), down_q.size()}; + model.file.reserve(gate_q.size() + up_q.size() + down_q.size()); + model.file.insert(model.file.end(), gate_q.begin(), gate_q.end()); + model.file.insert(model.file.end(), up_q.begin(), up_q.end()); + model.file.insert(model.file.end(), down_q.begin(), down_q.end()); + model.slot_bytes = model.regions.expert_bytes_gate + + model.regions.expert_bytes_up + + model.regions.expert_bytes_down; + return model; +} + std::vector cpu_reference( const std::vector & gate, const std::vector & up, @@ -289,9 +360,85 @@ void run_layout_case(ggml_backend_t backend, bool expert_major) { engine.destroy(); } +void run_mxfp4_padding_case(ggml_backend_t backend) { + std::vector gate; + std::vector up; + std::vector down; + fill_weights(gate, up, down); + std::vector gate_dequantized; + std::vector up_dequantized; + std::vector down_dequantized; + ModelBytes model = make_mxfp4_model_bytes( + gate, up, down, gate_dequantized, up_dequantized, + down_dequantized); + TempFile file(model.file); + + MoeHybridStorage storage; + storage.mmap_size = model.file.size(); + storage.mmap_fd = ::dup(file.fd); + STREAM_REQUIRE(storage.mmap_fd >= 0); + storage.layer_regions.push_back(model.regions); + + MoeStreamConfig config; + config.device_slots = 2; + config.graph_cache_entries = 4; + config.nvme.backend = MoeNvmeBackend::ThreadPool; + config.nvme.direct_io = MoeNvmeDirectMode::Disabled; + config.nvme.host_slots = 6; + + MoeHybridStreamEngine engine; + std::string error; + STREAM_REQUIRE(engine.init( + backend, model.slot_bytes, storage, config, &error)); + + MoeStreamExpertSpec spec; + spec.input_dim = kInput; + spec.intermediate_dim = kFf; + spec.output_dim = kOutput; + spec.gate_type = GGML_TYPE_MXFP4; + spec.up_type = GGML_TYPE_MXFP4; + spec.down_type = GGML_TYPE_MXFP4; + spec.gated_activation = MoeGatedActivation::Situ; + spec.gate_scale = 0.8f; + spec.up_scale = 1.1f; + spec.down_scale = 0.9f; + + std::vector input(static_cast(kTokens) * kInput); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = 0.12f * std::sin(0.07f * static_cast(i + 1)); + } + const int32_t ids[kTokens * kTopK] = {2, 0, 1, 2}; + const float weights[kTokens * kTopK] = {0.65f, 0.35f, 0.55f, 0.45f}; + MoeStreamRouteBatch batch; + batch.layer = 0; + batch.n_expert = kExperts; + batch.top_k = kTopK; + batch.n_tokens = kTokens; + batch.inputs = input.data(); + batch.selected_ids = ids; + batch.selected_weights = weights; + + const std::vector expected = cpu_reference( + gate_dequantized, up_dequantized, down_dequantized, + input, ids, weights); + std::vector actual; + STREAM_REQUIRE(eval_moe_streamed_experts( + engine, spec, batch, actual, &error)); + STREAM_REQUIRE(actual.size() == expected.size()); + for (size_t i = 0; i < actual.size(); ++i) { + const float tolerance = + 2.0e-4f + 2.0e-3f * std::fabs(expected[i]); + STREAM_REQUIRE(std::fabs(actual[i] - expected[i]) <= tolerance); + } + const MoeStreamComputeStats stats = engine.compute_stats(); + STREAM_REQUIRE(stats.graph_builds == 2); + STREAM_REQUIRE(stats.graph_launches == 3); + engine.destroy(); +} + } // namespace -TEST_CASE(MoeStreamComputeFixture, persistent_graph_matches_cpu_for_both_layouts) { +TEST_CASE(MoeStreamComputeFixture, persistent_graph_matches_cpu_and_padded_mxfp4) { int device = 0; if (const char * value = std::getenv("DFLASH_TEST_GPU")) { device = std::max(0, std::atoi(value)); @@ -307,5 +454,6 @@ TEST_CASE(MoeStreamComputeFixture, persistent_graph_matches_cpu_for_both_layouts } run_layout_case(backend, false); run_layout_case(backend, true); + run_mxfp4_padding_case(backend); ggml_backend_free(backend); } From bc4eb28ebc69ed9aeec1e804116f291d12725b99 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:28:37 +0200 Subject: [PATCH 07/20] ci(moe): run MXFP4 stream test on GPU runners --- .github/workflows/ci.yml | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3cfe3e446..350960916 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -174,6 +174,7 @@ jobs: -DCMAKE_BUILD_TYPE=Release cmake --build build \ --target test_flash_attn_sparse test_deepseek4_mmid_grouped_cuda \ + test_moe_stream_compute \ -j"$(nproc)" - name: Run flash-attn sparse kernel test on the 3090 @@ -184,6 +185,9 @@ jobs: - name: Run grouped MMID dispatch and parity test on the 3090 run: ./server/build/test_deepseek4_mmid_grouped_cuda + - name: Run compact-to-padded MXFP4 expert streaming test on the 3090 + run: ./server/build/test_moe_stream_compute + # Optional model-backed end-to-end smoke (real spec-decode on the 3090), # disabled by default because it builds dflash_server and lazy-loads the # ~16 GB Qwen3.6-27B target + draft (~1-2 min). The weights are already @@ -290,11 +294,12 @@ jobs: -DCMAKE_HIP_FLAGS=-DDFLASH_WAVE_SIZE=32 cmake --build "$RUNNER_TEMP/rocmfp-build" \ --target test_rocmfp4 test_rocmfpx test_rocmfp4_hip_tail test_rocmfpx_mmq \ - test_deepseek4_mmid_grouped_cuda test_recurrent_snapshot test_server_unit \ + test_deepseek4_mmid_grouped_cuda test_moe_stream_compute \ + test_recurrent_snapshot test_server_unit \ --parallel 8 ctest --test-dir "$RUNNER_TEMP/rocmfp-build" \ --output-on-failure \ - -R 'rocmfp4_reference|rocmfpx_reference|rocmfp4_hip_tail|rocmfpx_mmq|deepseek4_mmid_grouped_cuda|recurrent_snapshot|ChainRollbackPolicy' + -R 'rocmfp4_reference|rocmfpx_reference|rocmfp4_hip_tail|rocmfpx_mmq|deepseek4_mmid_grouped_cuda|test_moe_stream_compute|recurrent_snapshot|ChainRollbackPolicy' build-windows: name: Build Windows (MSVC + CUDA, library + server targets) From 1b25c9afe24a9cd47c27250fedf48b17d9deca14 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 31 Jul 2026 10:49:57 +0200 Subject: [PATCH 08/20] feat(kimi): stream routed experts from split GGUF --- server/docs/KIMI_K3_HETERO.md | 92 ++++-- server/docs/MOE_NVME_STREAMING.md | 35 ++ server/src/common/moe_hybrid_stream.cpp | 51 +-- server/src/common/moe_hybrid_stream.h | 3 + server/src/kimi_k3/kimi_k3_backend.cpp | 194 ++++++++++- server/src/kimi_k3/kimi_k3_backend.h | 7 + server/src/kimi_k3/kimi_k3_graph.cpp | 422 +++++++++++++++++++++++- server/src/kimi_k3/kimi_k3_internal.h | 22 +- server/src/kimi_k3/kimi_k3_loader.cpp | 388 +++++++++++++++++++--- server/test/smoke_kimi_k3_forward.cpp | 5 +- 10 files changed, 1089 insertions(+), 130 deletions(-) diff --git a/server/docs/KIMI_K3_HETERO.md b/server/docs/KIMI_K3_HETERO.md index 85ad15e3b..989c4997e 100644 --- a/server/docs/KIMI_K3_HETERO.md +++ b/server/docs/KIMI_K3_HETERO.md @@ -6,8 +6,8 @@ files and 594 GB. ## What is implemented -The common MoE data plane now has the Kimi-specific capabilities that were -missing without making the scheduler Kimi-specific: +Kimi K3 now uses the same model-neutral SSD data plane as DeepSeek, with a +small architecture adapter around it: - A tensor region carries a model-shard index. `MoeNvmeScheduler` can register and read any number of GGUF shard file descriptors with `io_uring`, including @@ -17,6 +17,18 @@ missing without making the scheduler Kimi-specific: its experts consume full-width hidden states. - The common streamed graph supports SiTU as well as SwiGLU and DS4's clamped SwiGLU. +- The native Kimi text backend executes KDA, absorbed MLA, Attention Residuals, + dense layer 0, latent routed MoE, shared experts, final projection, + tokenization, and sampling. Its own router IDs and weights are copied across + the activation-sized boundary into `eval_moe_streamed_experts`; the SSD + scheduler never substitutes or predicts a contributing expert. +- The loader accepts standard split GGUFs, even though only shard 1 contains + global metadata. It keeps routed gate/up/down stacks file-backed while + selectively allocating all other tensors on the target device. +- Device-cache sizing happens after resident weights and recurrent/KV state + are allocated. It reserves the larger of 2 GiB or 5% of device memory and is + capped by the actual routed pool, so a tiny model cannot accidentally request + a model-sized cache. - `bench_kimi_k3_hetero` runs Kimi's exact IQ1_S routed-expert geometry through the same model-neutral persistent evaluator used by production adapters: 896 experts, top-16, 92 MoE layers, @@ -29,11 +41,35 @@ matches both tensor-major and expert-major streamed GPU execution against a CPU oracle on gfx1151 and gfx1201. Existing single-file DS4 descriptors remain source index zero and need no model-specific change. -This is not yet a complete Kimi K3 backend. KDA/MLA, Attention Residuals, the -vision encoder, the latent projections around the routed core, tokenizer, and -sampling still need a model adapter. The current upstream llama.cpp Kimi K3 -text-model implementation is also not merged, so it should be treated as a -reference implementation rather than a stable dependency. +The SSD text path is implemented, but two qualification boundaries remain: + +- The backend is correctness-first and token-sequential. Its per-layer graph + boundaries are not yet fused/captured for full-model speed. +- Kimi currently places resident text tensors and streamed expert compute on + one selected GPU. Strix-only is supported directly. Splitting Kimi's dense + tensors onto R9700 while Strix owns streamed experts is a later performance + adapter, not a requirement for SSD capacity correctness. + +The vision encoder is out of scope for this text-only path. + +## End-to-end split-GGUF qualification, 2026-07-31 + +The released `inference-optimization/Kimi-K3-0.40B-MXFP4` architecture fixture +was converted to two standard GGUF shards. The routed tensors remained native +MXFP4. On Strix Halo, the same prompt was run once with all experts resident +and once with all routed stacks file-backed: + +```text +prompt IDs: 18805 308 799 5624 12524 +output IDs: 318 57195 11 1459 387 1495 2189 261 +text: According to all known laws of aviation, there is no way a +``` + +The token IDs matched exactly. The automatic Linux path selected `io_uring`, +read all selected experts from the correct GGUF shards, and reported 168 +expert launches, one graph build, 167 graph-cache hits, and zero I/O errors. +With a deliberately tiny 1 MiB device cache it moved 0.033 GiB; this forces +evictions and proves the result is not an all-resident accident. ## Exact routed-weight demand @@ -90,41 +126,49 @@ a cache only helps in proportion to its share of the 495 GiB routed pool. The machine has about 125.08 GiB of system/UMA memory plus 31.86 GiB on the R9700, or 156.94 GiB of unique physical weight capacity before runtime -reserves. A simple placement is: +reserves. The implemented capacity-safe starting point is: -1. R9700: attention/KDA/MLA and other dense matrices that fit its 32 GiB. -2. Strix/system memory: remaining non-routed weights, shared experts, latent +1. Strix/system memory: all non-routed text weights, shared experts, latent projections, recurrent/KV state, workspace, and the routed-expert cache. -3. NVMe: all routed expert stacks, with actual route misses read directly into +2. NVMe: all routed expert stacks, with actual route misses read directly into pinned slots and evaluated on Strix. +This works unchanged on a Strix-only machine. On a full Lucebox the R9700 is +currently unused by the Kimi adapter; moving dense KDA/MLA work there is the +next throughput optimization. The common SSD runtime already supports a +different compute owner, but Kimi still needs device-aware resident allocation +and cache ownership before that placement is correct end to end. + After approximately 57.94 GiB of non-routed weights plus OS, workspace, and a moderate context reserve, roughly 70-85 GiB may remain for routed experts. Under a uniform balanced-routing assumption this covers about 14-17% of the routed pool. At the measured 3.804 GiB/s, the storage-only ceiling is then -approximately 0.50-0.51 token/s. Full inference will be lower unless dense -R9700 work overlaps almost completely with Strix expert service. +approximately 0.50-0.51 token/s. Full Strix-only inference will be lower +because dense/recurrent work shares the same device; a later R9700 split may +recover part of that gap through overlap. So the honest expectation for this quant is **roughly one token every two to three seconds**, not interactive multi-token-per-second generation. Real router locality can move that estimate; only a route trace from the real model can establish it. -## Next end-to-end milestone +## Next full-model milestone -The next useful step is one narrow Kimi adapter, not another generic cache: +The implementation no longer needs another generic cache or a second Kimi +router. Full-scale qualification requires: -1. Import the correctness-first text graph from the upstream Kimi K3 work. -2. Populate shard-indexed expert regions directly from the GGUF tensor table. -3. Place dense attention on R9700 and the latent/shared/routed MoE path on - Strix; keep only activation-sized transfers at the boundary. -4. Record real `(layer, expert)` routes on a calibration prompt suite and let +1. Stage all 14 IQ1_S shards and run a short token-for-token comparison against + the upstream Kimi implementation. +2. Record real `(layer, expert)` routes on a calibration prompt suite and let the existing placement planner allocate the measured best cache under the chosen context budget. -5. Compare end-to-end output and token rate against ordinary llama.cpp +3. Compare end-to-end output and token rate against ordinary llama.cpp CPU/GPU offload. +4. Only after that baseline, place dense attention on R9700 and the + latent/shared/routed MoE path on Strix, retaining activation-sized transfers + at the boundary. The full 594 GB model cannot currently be staged on the qualification box, -which has substantially less free SSD space. It needs at least about 650 GB of -safe free space for all shards plus conversion/logging headroom; existing user -models should not be deleted implicitly. +which currently has about 513 GB free. It needs at least about 650 GB of safe +free space for all shards plus logging/headroom; existing user models should +not be deleted implicitly. diff --git a/server/docs/MOE_NVME_STREAMING.md b/server/docs/MOE_NVME_STREAMING.md index 5b8ac960b..52e17f250 100644 --- a/server/docs/MOE_NVME_STREAMING.md +++ b/server/docs/MOE_NVME_STREAMING.md @@ -118,6 +118,30 @@ model that genuinely exceeds UMA, `auto` chooses the partial placement itself. `DFLASH_EXPERT_BUDGET_MB` can additionally cap static routed-weight residency; `DFLASH_MOE_NVME_DEVICE_CACHE_MB` can override the adaptive cache budget. +## Kimi K3 activation + +Pass the first file of a standard split GGUF and select the Strix device. +Kimi's routed gate/up/down stacks remain file-backed by default; the loader +allocates only dense, shared, routing, latent-projection, and cache tensors. + +```bash +export DFLASH_MOE_NVME_BACKEND=auto + +./build-hip/dflash_server \ + /path/to/Kimi-K3-UD-IQ1_S-00001-of-00014.gguf \ + --target-device hip:0 --max-ctx 8192 +``` + +No DeepSeek MoE-TP variables are required. Cache memory is chosen only after +the resident model and recurrent/KV state have been allocated. The automatic +budget reserves the larger of 2 GiB or 5% of device memory and never exceeds +the complete routed pool. `DFLASH_MOE_NVME_DEVICE_CACHE_MB` remains an explicit +override, also capped by the routed pool. + +Kimi's native router remains authoritative. The current text backend is +correctness-first and sequential; multi-device dense placement, captured +per-layer graphs, and the vision tower are separate optimizations. + ## Tuning and diagnostics Defaults are intentionally small: eight pinned host slots, four fallback I/O @@ -147,6 +171,10 @@ targets `test_moe_nvme_scheduler`, `bench_moe_nvme_io`, and `bench_moe_nvme_pipeline` test scheduling, raw storage, and the complete SSD-to-GPU path. Benchmarks are read-only. +The external `smoke_kimi_k3_forward` target accepts +`[stream_experts=0|1]`. This is an A/B oracle for small Kimi fixtures; production +Kimi uses `1`. + ## Qualification result (2026-07-30) On the Lucebox P310 and a realistic 24 MiB expert working set: @@ -186,6 +214,13 @@ took 3.816 and 3.599 seconds on the same run shape, so persistence improved this deliberately cold, storage-bound case by 1.4-2.7%; its larger value is removing thousands of allocations when more of the route set is warm. +The native Kimi path was qualified with a real two-shard 0.40B MXFP4 +architecture fixture. Resident and streamed execution produced the same eight +greedy output tokens. The `io_uring` run completed 168 selected-expert +launches with one graph build, 167 graph-cache hits, and zero I/O errors. A +1 MiB device cache forced 163 evictions, demonstrating that the result came +through the split-GGUF SSD path rather than accidental full residency. + ## Research lineage and next optimization The bounded priority/cache design follows the lessons of MoE-Infinity and diff --git a/server/src/common/moe_hybrid_stream.cpp b/server/src/common/moe_hybrid_stream.cpp index ec3b4cc3a..1f78ba594 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -808,6 +808,13 @@ MoeHybridStreamEngine & MoeHybridStreamEngine::operator=(MoeHybridStreamEngine & bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, size_t max_expert_bytes, std::string * err) { + return init(gpu_backend, max_expert_bytes, MoeStreamConfig::from_env(), err); +} + +bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, + size_t max_expert_bytes, + const MoeStreamConfig & config, + std::string * err) { destroy(); if (!gpu_backend || max_expert_bytes == 0) { if (err) *err = "invalid arguments to stream engine init"; @@ -827,7 +834,8 @@ bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, size_t max_expert_b return false; } runtime->max_expert_bytes = max_expert_bytes; - runtime->config = MoeStreamConfig::from_env(); + runtime->config = config; + runtime->config.device_slots = std::max(2, runtime->config.device_slots); runtime->io.reset(new (std::nothrow) MoeNvmeScheduler); if (!runtime->io) { if (err) *err = "failed to allocate SSD scheduler"; @@ -864,46 +872,7 @@ bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, size_t max_expert_b const MoeHybridStorage & storage, const MoeStreamConfig & config, std::string * err) { - destroy(); - if (!gpu_backend || max_expert_bytes == 0) { - if (err) *err = "invalid arguments to stream engine init"; - return false; - } - - std::unique_ptr runtime(new (std::nothrow) Runtime); - if (!runtime) { - if (err) *err = "failed to allocate stream runtime"; - return false; - } - runtime->backend = gpu_backend; - runtime->device = backend_device_index(gpu_backend); - ScopedGpuDevice device_scope(runtime->device); - if (!device_scope.ready()) { - if (err) *err = "failed to resolve/select SSD stream GPU"; - return false; - } - runtime->max_expert_bytes = max_expert_bytes; - runtime->config = config; - runtime->config.device_slots = std::max(2, runtime->config.device_slots); - runtime->io.reset(new (std::nothrow) MoeNvmeScheduler); - if (!runtime->io || - !runtime->io->init(runtime->config.nvme, max_expert_bytes, - pinned_allocate, pinned_free, nullptr, err)) { - return false; - } - - cudaError_t gpu_err = cudaStreamCreate(&runtime->transfer_stream); - if (gpu_err != cudaSuccess) { - if (err) *err = std::string("failed to create SSD transfer stream: ") + - cudaGetErrorString(gpu_err); - return false; - } - if (!allocate_device_cache(*runtime, err)) { - runtime_ = std::move(runtime); - destroy(); - return false; - } - runtime_ = std::move(runtime); + if (!init(gpu_backend, max_expert_bytes, config, err)) return false; if (!bind_storage(storage, err)) { destroy(); return false; diff --git a/server/src/common/moe_hybrid_stream.h b/server/src/common/moe_hybrid_stream.h index e267c6867..350606464 100644 --- a/server/src/common/moe_hybrid_stream.h +++ b/server/src/common/moe_hybrid_stream.h @@ -117,6 +117,9 @@ class MoeHybridStreamEngine { // available instead of relying only on mmap page faults. bool init(ggml_backend_t gpu_backend, size_t max_expert_bytes, std::string * err = nullptr); + bool init(ggml_backend_t gpu_backend, size_t max_expert_bytes, + const MoeStreamConfig & config, + std::string * err = nullptr); bool init(ggml_backend_t gpu_backend, size_t max_expert_bytes, const MoeHybridStorage & storage, std::string * err = nullptr); diff --git a/server/src/kimi_k3/kimi_k3_backend.cpp b/server/src/kimi_k3/kimi_k3_backend.cpp index bf9873273..269862697 100644 --- a/server/src/kimi_k3/kimi_k3_backend.cpp +++ b/server/src/kimi_k3/kimi_k3_backend.cpp @@ -6,11 +6,25 @@ #include "ggml-cuda.h" #include +#include #include #include +#include +#include +#include #include #include +#if defined(_WIN32) +#include +#include +#include +#else +#include +#include +#include +#endif + namespace dflash::common { KimiK3Backend::KimiK3Backend(const KimiK3BackendConfig & cfg) : cfg_(cfg) {} @@ -19,6 +33,163 @@ KimiK3Backend::~KimiK3Backend() { shutdown(); } +bool KimiK3Backend::init_streaming() { + if (!weights_.routed_experts_streamed || + weights_.streamed_layer_regions.empty() || + weights_.max_streamed_expert_bytes == 0) { + std::fprintf(stderr, + "[kimi-k3] routed expert streaming metadata is incomplete\n"); + return false; + } + + MoeStreamConfig stream_config = MoeStreamConfig::from_env(); + size_t routed_pool_bytes = 0; + for (const LayerExpertRegions & regions : + weights_.streamed_layer_regions) { + size_t bytes_per_expert = 0; + bool component_overflow = false; + for (size_t component : { + regions.expert_bytes_gate, regions.expert_bytes_up, + regions.expert_bytes_down, regions.expert_bytes_gate_up}) { + if (component > + std::numeric_limits::max() - bytes_per_expert) { + component_overflow = true; + break; + } + bytes_per_expert += component; + } + if (component_overflow) { + routed_pool_bytes = std::numeric_limits::max(); + break; + } + if (bytes_per_expert > + (std::numeric_limits::max() - routed_pool_bytes) / + static_cast(weights_.n_expert)) { + routed_pool_bytes = std::numeric_limits::max(); + break; + } + routed_pool_bytes += + bytes_per_expert * static_cast(weights_.n_expert); + } + if (!std::getenv("DFLASH_MOE_NVME_DEVICE_CACHE_MB")) { + size_t free_bytes = 0; + size_t total_bytes = 0; + ggml_backend_cuda_get_device_memory( + cfg_.device.primary_gpu(), &free_bytes, &total_bytes); + const size_t gib = 1024ULL * 1024ULL * 1024ULL; + const size_t reserve = std::max(2 * gib, total_bytes / 20); + stream_config.device_cache_bytes = + free_bytes > reserve + ? std::min(free_bytes - reserve, routed_pool_bytes) + : 0; + std::fprintf(stderr, + "[kimi-k3] streamed expert cache: free=%.2f GiB reserve=%.2f GiB " + "pool=%.2f GiB cache=%.2f GiB\n", + static_cast(free_bytes) / gib, + static_cast(reserve) / gib, + static_cast(routed_pool_bytes) / gib, + static_cast(stream_config.device_cache_bytes) / gib); + } else { + stream_config.device_cache_bytes = + std::min(stream_config.device_cache_bytes, routed_pool_bytes); + } + std::string error; + if (!stream_engine_.init( + backend_, weights_.max_streamed_expert_bytes, + stream_config, &error)) { + std::fprintf(stderr, + "[kimi-k3] stream engine initialization failed: %s\n", + error.c_str()); + return false; + } + + std::vector descriptors; + std::vector sources; + descriptors.reserve(weights_.shard_paths.size()); + sources.reserve(weights_.shard_paths.size()); + for (const std::string & shard : weights_.shard_paths) { +#if defined(_WIN32) + const int fd = ::_open(shard.c_str(), _O_RDONLY | _O_BINARY); +#else + const int fd = ::open(shard.c_str(), O_RDONLY | O_CLOEXEC); +#endif + if (fd < 0) { + std::fprintf(stderr, + "[kimi-k3] cannot open expert shard %s: %s\n", + shard.c_str(), std::strerror(errno)); + for (int opened : descriptors) { +#if defined(_WIN32) + ::_close(opened); +#else + ::close(opened); +#endif + } + stream_engine_.destroy(); + return false; + } + uint64_t shard_bytes = 0; +#if defined(_WIN32) + struct _stat64 stat_buffer {}; + if (::_fstat64(fd, &stat_buffer) == 0 && stat_buffer.st_size > 0) { + shard_bytes = static_cast(stat_buffer.st_size); + } +#else + struct stat stat_buffer {}; + if (::fstat(fd, &stat_buffer) == 0 && stat_buffer.st_size > 0) { + shard_bytes = static_cast(stat_buffer.st_size); + } +#endif + if (shard_bytes == 0 || + shard_bytes > std::numeric_limits::max()) { + std::fprintf(stderr, + "[kimi-k3] cannot determine expert shard size: %s\n", + shard.c_str()); +#if defined(_WIN32) + ::_close(fd); +#else + ::close(fd); +#endif + for (int opened : descriptors) { +#if defined(_WIN32) + ::_close(opened); +#else + ::close(opened); +#endif + } + stream_engine_.destroy(); + return false; + } + descriptors.push_back(fd); + sources.push_back({ + nullptr, static_cast(shard_bytes), fd}); + } + const bool bound = stream_engine_.bind_sources( + sources, weights_.streamed_layer_regions, &error); + for (int fd : descriptors) { +#if defined(_WIN32) + ::_close(fd); +#else + ::close(fd); +#endif + } + if (!bound) { + std::fprintf(stderr, + "[kimi-k3] stream source binding failed: %s\n", + error.c_str()); + stream_engine_.destroy(); + return false; + } + std::fprintf(stderr, + "[kimi-k3] routed experts file-backed: shards=%zu layers=%zu " + "io=%s cache=%.2f GiB\n", + weights_.shard_paths.size(), + weights_.streamed_layer_regions.size(), + stream_engine_.io_backend_name(), + static_cast(stream_engine_.device_cache_bytes()) / + (1024.0 * 1024.0 * 1024.0)); + return true; +} + bool KimiK3Backend::init() { if (!cfg_.model_path) { std::fprintf(stderr, "[kimi-k3] model path is null\n"); @@ -30,7 +201,9 @@ bool KimiK3Backend::init() { cfg_.device.primary_gpu()); return false; } - if (!load_kimi_k3_gguf(cfg_.model_path, backend_, weights_)) { + if (!load_kimi_k3_gguf( + cfg_.model_path, backend_, weights_, + cfg_.stream_routed_experts)) { std::fprintf(stderr, "[kimi-k3] model load failed: %s\n", dflash27b_last_error()); return false; @@ -41,10 +214,12 @@ bool KimiK3Backend::init() { max_ctx); return false; } + if (weights_.routed_experts_streamed && !init_streaming()) return false; std::fprintf(stderr, "[kimi-k3] native backend ready on device %d (max_ctx=%d, " - "correctness-first sequential prefill)\n", - cfg_.device.primary_gpu(), max_ctx); + "experts=%s, correctness-first sequential prefill)\n", + cfg_.device.primary_gpu(), max_ctx, + weights_.routed_experts_streamed ? "nvme" : "resident"); std::fflush(stderr); return true; } @@ -60,6 +235,7 @@ void KimiK3Backend::print_ready_banner() const { bool KimiK3Backend::park(ParkTarget target) { if (!park_target_includes_target_model(target)) return false; if (!parked_) { + stream_engine_.destroy(); free_kimi_k3_weights(weights_); parked_ = true; } @@ -69,7 +245,12 @@ bool KimiK3Backend::park(ParkTarget target) { bool KimiK3Backend::unpark(ParkTarget target) { if (!park_target_includes_target_model(target)) return false; if (parked_) { - if (!load_kimi_k3_gguf(cfg_.model_path, backend_, weights_)) return false; + if (!load_kimi_k3_gguf( + cfg_.model_path, backend_, weights_, + cfg_.stream_routed_experts) || + (weights_.routed_experts_streamed && !init_streaming())) { + return false; + } parked_ = false; } return true; @@ -114,7 +295,7 @@ GenerateResult KimiK3Backend::generate_impl(const GenerateRequest & req, const auto prefill_begin = std::chrono::steady_clock::now(); for (size_t i = 0; i < req.prompt.size(); ++i) { if (!kimi_k3_step(backend_, weights_, cache_, req.prompt[i], - static_cast(i), logits)) { + static_cast(i), logits, &stream_engine_)) { result.fail(GenerateErrorCode::PrefillFailed, dflash27b_last_error()); out_io.emit(-1); @@ -156,7 +337,7 @@ GenerateResult KimiK3Backend::generate_impl(const GenerateRequest & req, if (out_io.cancelled || next == weights_.eos_token_id) break; if (i + 1 < req.n_gen) { if (!kimi_k3_step(backend_, weights_, cache_, next, - cache_.cur_pos, logits)) { + cache_.cur_pos, logits, &stream_engine_)) { result.fail(GenerateErrorCode::DecodeFailed, dflash27b_last_error()); out_io.emit(-1); @@ -209,6 +390,7 @@ bool KimiK3Backend::handle_compress(const std::string & line, } void KimiK3Backend::shutdown() { + stream_engine_.destroy(); free_kimi_k3_cache(cache_); free_kimi_k3_weights(weights_); if (backend_) { diff --git a/server/src/kimi_k3/kimi_k3_backend.h b/server/src/kimi_k3/kimi_k3_backend.h index d15d37416..33b57b000 100644 --- a/server/src/kimi_k3/kimi_k3_backend.h +++ b/server/src/kimi_k3/kimi_k3_backend.h @@ -1,6 +1,7 @@ #pragma once #include "common/model_backend.h" +#include "common/moe_hybrid_stream.h" #include "kimi_k3_internal.h" #include "placement/placement_config.h" @@ -13,6 +14,9 @@ struct KimiK3BackendConfig { const char * model_path = nullptr; DevicePlacement device; int stream_fd = -1; + // Production Kimi uses file-backed routed experts. The resident mode is + // retained as a deterministic oracle for small architecture fixtures. + bool stream_routed_experts = true; }; class KimiK3Backend final : public ModelBackend { @@ -44,6 +48,8 @@ class KimiK3Backend final : public ModelBackend { void shutdown() override; private: + bool init_streaming(); + int32_t choose_token(const std::vector & logits, const SamplerCfg & sampler, const std::vector & history); @@ -52,6 +58,7 @@ class KimiK3Backend final : public ModelBackend { ggml_backend_t backend_ = nullptr; KimiK3Weights weights_; KimiK3Cache cache_; + MoeHybridStreamEngine stream_engine_; bool parked_ = false; std::mt19937_64 rng_{std::random_device{}()}; }; diff --git a/server/src/kimi_k3/kimi_k3_graph.cpp b/server/src/kimi_k3/kimi_k3_graph.cpp index b5e601447..da6e0810b 100644 --- a/server/src/kimi_k3/kimi_k3_graph.cpp +++ b/server/src/kimi_k3/kimi_k3_graph.cpp @@ -1,5 +1,6 @@ #include "kimi_k3_internal.h" +#include "common/moe_hybrid_stream.h" #include "common/moe_router_graph.h" #include "internal.h" @@ -279,15 +280,12 @@ ggml_tensor * build_mla(ggml_context * ctx, return ggml_mul_mat(ctx, layer.wo, out); } -ggml_tensor * build_latent_moe(ggml_context * ctx, - ggml_cgraph * graph, - const KimiK3Weights & w, - const KimiK3Layer & layer, - ggml_tensor * cur) { - ggml_tensor * identity = cur; - ggml_tensor * routed_in = ggml_mul_mat(ctx, layer.ffn_routed_down, cur); - ggml_tensor * logits = ggml_mul_mat(ctx, layer.ffn_gate_inp, identity); - +TopKMoeRouterResult build_kimi_router(ggml_context * ctx, + ggml_cgraph * graph, + const KimiK3Weights & w, + const KimiK3Layer & layer, + ggml_tensor * cur) { + ggml_tensor * logits = ggml_mul_mat(ctx, layer.ffn_gate_inp, cur); TopKMoeRouterResult router; if (w.expert_gating_func == 2) { router = build_sigmoid_topk_moe_router(ctx, graph, logits, @@ -311,6 +309,18 @@ ggml_tensor * build_latent_moe(ggml_context * ctx, router.weights_2d = weights; router.weights_3d = ggml_reshape_3d(ctx, weights, 1, w.n_expert_used, 1); } + return router; +} + +ggml_tensor * build_latent_moe(ggml_context * ctx, + ggml_cgraph * graph, + const KimiK3Weights & w, + const KimiK3Layer & layer, + ggml_tensor * cur) { + ggml_tensor * identity = cur; + ggml_tensor * routed_in = ggml_mul_mat(ctx, layer.ffn_routed_down, cur); + TopKMoeRouterResult router = + build_kimi_router(ctx, graph, w, layer, identity); ggml_tensor * routed_3d = ggml_reshape_3d(ctx, routed_in, w.n_expert_latent, 1, 1); @@ -340,6 +350,387 @@ ggml_tensor * build_latent_moe(ggml_context * ctx, return ggml_add(ctx, moe, shared); } +struct GraphInput { + ggml_tensor * tensor = nullptr; + const void * data = nullptr; + size_t bytes = 0; +}; + +struct GraphOutput { + ggml_tensor * tensor = nullptr; + void * data = nullptr; + size_t bytes = 0; +}; + +bool run_host_boundary_graph(ggml_backend_t backend, + ggml_context * ctx, + ggml_cgraph * graph, + const std::vector & inputs, + const std::vector & outputs, + const char * phase) { + for (const GraphOutput & output : outputs) { + if (!output.tensor || !output.data || output.bytes == 0) { + set_last_error(std::string("Kimi-K3 ") + phase + + ": invalid graph output"); + return false; + } + ggml_set_output(output.tensor); + ggml_build_forward_expand(graph, output.tensor); + } + ggml_gallocr_t allocator = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + if (!allocator || !ggml_gallocr_alloc_graph(allocator, graph)) { + set_last_error(std::string("Kimi-K3 ") + phase + + ": graph allocation failed"); + if (allocator) ggml_gallocr_free(allocator); + return false; + } + for (const GraphInput & input : inputs) { + if (!input.tensor || !input.data || input.bytes == 0) { + set_last_error(std::string("Kimi-K3 ") + phase + + ": invalid graph input"); + ggml_gallocr_free(allocator); + return false; + } + ggml_backend_tensor_set( + input.tensor, input.data, 0, input.bytes); + } + const ggml_status status = + ggml_backend_graph_compute(backend, graph); + if (status != GGML_STATUS_SUCCESS) { + set_last_error(std::string("Kimi-K3 ") + phase + + ": graph compute failed with status " + + std::to_string(static_cast(status))); + ggml_gallocr_free(allocator); + return false; + } + for (const GraphOutput & output : outputs) { + ggml_backend_tensor_get( + output.tensor, output.data, 0, output.bytes); + } + ggml_gallocr_free(allocator); + return true; +} + +void populate_attn_res_bank( + ggml_context * ctx, + const KimiK3Weights & w, + const std::vector> & host_checkpoints, + AttnResBank & bank, + std::vector & inputs) { + bank.ctx = ctx; + bank.eps = w.rms_eps; + bank.n_embd = w.n_embd; + for (const std::vector & checkpoint : host_checkpoints) { + ggml_tensor * tensor = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, w.n_embd, 1); + ggml_set_input(tensor); + inputs.push_back({ + tensor, checkpoint.data(), + checkpoint.size() * sizeof(float)}); + bank.push(tensor); + } +} + +ggml_context * new_kimi_step_context() { + ggml_init_params params{}; + params.mem_size = 64ull * 1024ull * 1024ull; + params.no_alloc = true; + return ggml_init(params); +} + +bool streamed_kimi_k3_step( + ggml_backend_t backend, + const KimiK3Weights & w, + KimiK3Cache & cache, + int32_t token, + int position, + std::vector & logits, + MoeHybridStreamEngine & stream_engine) { + std::vector hidden(static_cast(w.n_embd)); + + { + ggml_context * ctx = new_kimi_step_context(); + if (!ctx) { + set_last_error("Kimi-K3 embedding: context allocation failed"); + return false; + } + ggml_cgraph * graph = + ggml_new_graph_custom(ctx, 1024, false); + ggml_tensor * ids = + ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_set_input(ids); + ggml_tensor * embedding = + ggml_get_rows(ctx, w.tok_embd, ids); + const bool ok = run_host_boundary_graph( + backend, ctx, graph, + {{ids, &token, sizeof(token)}}, + {{embedding, hidden.data(), + hidden.size() * sizeof(float)}}, + "embedding"); + ggml_free(ctx); + if (!ok) return false; + } + + std::vector> checkpoints; + checkpoints.reserve( + static_cast( + (w.n_layer + w.attn_res_block_size - 1) / + w.attn_res_block_size)); + + for (int il = 0; il < w.n_layer; ++il) { + const KimiK3Layer & layer = + w.layers[static_cast(il)]; + KimiK3LayerCache & layer_cache = + cache.layers[static_cast(il)]; + const bool banked = + il % w.attn_res_block_size == 0; + const std::vector checkpoint_value = hidden; + + ggml_context * ctx = new_kimi_step_context(); + if (!ctx) { + set_last_error("Kimi-K3 layer: context allocation failed"); + return false; + } + ggml_cgraph * graph = + ggml_new_graph_custom(ctx, 32768, false); + std::vector inputs; + ggml_tensor * hidden_in = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, w.n_embd, 1); + ggml_set_input(hidden_in); + inputs.push_back({ + hidden_in, hidden.data(), + hidden.size() * sizeof(float)}); + + AttnResBank residuals; + populate_attn_res_bank( + ctx, w, checkpoints, residuals, inputs); + ggml_tensor * prefix = hidden_in; + ggml_tensor * cur = + residuals.mix(prefix, layer.attn_res_score); + if (banked) residuals.push(prefix); + + cur = rms_norm( + ctx, cur, layer.attn_norm, w.rms_eps); + cur = layer.recurrent + ? build_kda( + ctx, graph, w, layer, layer_cache, cur) + : build_mla( + ctx, graph, w, layer, layer_cache, + cur, position); + prefix = banked + ? cur : ggml_add(ctx, prefix, cur); + cur = residuals.mix(prefix, layer.ffn_res_score); + cur = rms_norm( + ctx, cur, layer.ffn_norm, w.rms_eps); + + if (il < w.n_dense_lead) { + ggml_tensor * gate = + ggml_mul_mat(ctx, layer.ffn_gate, cur); + ggml_tensor * up = + ggml_mul_mat(ctx, layer.ffn_up, cur); + ggml_tensor * dense = situ( + ctx, gate, up, + w.situ_beta, w.situ_linear_beta); + dense = ggml_mul_mat( + ctx, layer.ffn_down, dense); + ggml_tensor * hidden_out = + ggml_add(ctx, prefix, dense); + std::vector next_hidden( + static_cast(w.n_embd)); + const bool ok = run_host_boundary_graph( + backend, ctx, graph, inputs, + {{hidden_out, next_hidden.data(), + next_hidden.size() * sizeof(float)}}, + "dense layer"); + ggml_free(ctx); + if (!ok) return false; + if (banked) checkpoints.push_back(checkpoint_value); + hidden.swap(next_hidden); + continue; + } + + ggml_tensor * routed_in = + ggml_mul_mat(ctx, layer.ffn_routed_down, cur); + TopKMoeRouterResult router = + build_kimi_router(ctx, graph, w, layer, cur); + // argsort_top_k returns a strided view of the full argsort result. + // Materialize the tiny host-boundary tensors so the graph allocator + // cannot recycle their backing storage before the readback. + ggml_tensor * selected_out = + ggml_cont(ctx, router.selected); + ggml_tensor * route_weights_out = + ggml_cont(ctx, router.weights_2d); + ggml_tensor * shared_gate = + ggml_mul_mat(ctx, layer.ffn_gate_shexp, cur); + ggml_tensor * shared_up = + ggml_mul_mat(ctx, layer.ffn_up_shexp, cur); + ggml_tensor * shared = situ( + ctx, shared_gate, shared_up, + w.situ_beta, w.situ_linear_beta); + shared = ggml_mul_mat( + ctx, layer.ffn_down_shexp, shared); + + std::vector prefix_host( + static_cast(w.n_embd)); + std::vector routed_input_host( + static_cast(w.n_expert_latent)); + std::vector selected( + static_cast(w.n_expert_used)); + std::vector route_weights( + static_cast(w.n_expert_used)); + std::vector shared_host( + static_cast(w.n_embd)); + const bool prep_ok = run_host_boundary_graph( + backend, ctx, graph, inputs, + { + {prefix, prefix_host.data(), + prefix_host.size() * sizeof(float)}, + {routed_in, routed_input_host.data(), + routed_input_host.size() * sizeof(float)}, + {selected_out, selected.data(), + selected.size() * sizeof(int32_t)}, + {route_weights_out, route_weights.data(), + route_weights.size() * sizeof(float)}, + {shared, shared_host.data(), + shared_host.size() * sizeof(float)}, + }, + "routed layer preparation"); + ggml_free(ctx); + if (!prep_ok) return false; + if (banked) checkpoints.push_back(checkpoint_value); + for (size_t route = 0; route < selected.size(); ++route) { + if (selected[route] < 0 || selected[route] >= w.n_expert) { + set_last_error( + "Kimi-K3 routed layer " + std::to_string(il) + + ": native router returned invalid expert " + + std::to_string(selected[route]) + " at route " + + std::to_string(route)); + return false; + } + if (!std::isfinite(route_weights[route])) { + set_last_error( + "Kimi-K3 routed layer " + std::to_string(il) + + ": native router returned a non-finite weight"); + return false; + } + } + + MoeStreamExpertSpec spec; + spec.input_dim = w.n_expert_latent; + spec.intermediate_dim = w.n_ff_exp; + spec.output_dim = w.n_expert_latent; + spec.gate_type = layer.ffn_gate_exps->type; + spec.up_type = layer.ffn_up_exps->type; + spec.down_type = layer.ffn_down_exps->type; + spec.gated_activation = MoeGatedActivation::Situ; + spec.situ_beta = w.situ_beta; + spec.situ_linear_beta = w.situ_linear_beta; + + MoeStreamRouteBatch route_batch; + route_batch.layer = il - w.n_dense_lead; + route_batch.n_expert = w.n_expert; + route_batch.top_k = w.n_expert_used; + route_batch.n_tokens = 1; + route_batch.inputs = routed_input_host.data(); + route_batch.selected_ids = selected.data(); + route_batch.selected_weights = route_weights.data(); + std::vector routed_output; + std::string stream_error; + if (!eval_moe_streamed_experts( + stream_engine, spec, route_batch, + routed_output, &stream_error)) { + set_last_error( + "Kimi-K3 routed layer " + + std::to_string(il) + + ": streamed expert evaluation failed: " + + stream_error); + return false; + } + + ctx = new_kimi_step_context(); + if (!ctx) { + set_last_error( + "Kimi-K3 routed layer join: context allocation failed"); + return false; + } + graph = ggml_new_graph_custom(ctx, 4096, false); + ggml_tensor * prefix_in = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, w.n_embd, 1); + ggml_tensor * routed_out_in = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, w.n_expert_latent, 1); + ggml_tensor * shared_in = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, w.n_embd, 1); + ggml_set_input(prefix_in); + ggml_set_input(routed_out_in); + ggml_set_input(shared_in); + ggml_tensor * routed = routed_out_in; + if (layer.ffn_routed_norm) { + routed = rms_norm( + ctx, routed, layer.ffn_routed_norm, w.rms_eps); + } + routed = ggml_mul_mat( + ctx, layer.ffn_routed_up, routed); + ggml_tensor * moe_shared = + ggml_add(ctx, routed, shared_in); + ggml_tensor * hidden_out = + ggml_add(ctx, prefix_in, moe_shared); + std::vector next_hidden( + static_cast(w.n_embd)); + const bool join_ok = run_host_boundary_graph( + backend, ctx, graph, + { + {prefix_in, prefix_host.data(), + prefix_host.size() * sizeof(float)}, + {routed_out_in, routed_output.data(), + routed_output.size() * sizeof(float)}, + {shared_in, shared_host.data(), + shared_host.size() * sizeof(float)}, + }, + {{hidden_out, next_hidden.data(), + next_hidden.size() * sizeof(float)}}, + "routed layer join"); + ggml_free(ctx); + if (!join_ok) return false; + hidden.swap(next_hidden); + } + + ggml_context * ctx = new_kimi_step_context(); + if (!ctx) { + set_last_error("Kimi-K3 output: context allocation failed"); + return false; + } + ggml_cgraph * graph = + ggml_new_graph_custom(ctx, 8192, false); + std::vector inputs; + ggml_tensor * hidden_in = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, w.n_embd, 1); + ggml_set_input(hidden_in); + inputs.push_back({ + hidden_in, hidden.data(), + hidden.size() * sizeof(float)}); + AttnResBank residuals; + populate_attn_res_bank( + ctx, w, checkpoints, residuals, inputs); + ggml_tensor * output_hidden = + residuals.mix(hidden_in, w.output_res_score); + output_hidden = rms_norm( + ctx, output_hidden, w.output_norm, w.rms_eps); + ggml_tensor * output = + ggml_mul_mat(ctx, w.output, output_hidden); + logits.resize(static_cast(w.n_vocab)); + const bool output_ok = run_host_boundary_graph( + backend, ctx, graph, inputs, + {{output, logits.data(), + logits.size() * sizeof(float)}}, + "output"); + ggml_free(ctx); + if (!output_ok) return false; + + cache.cur_pos = position + 1; + return true; +} + } // namespace bool create_kimi_k3_cache(ggml_backend_t backend, @@ -405,13 +796,24 @@ bool kimi_k3_step(ggml_backend_t backend, KimiK3Cache & cache, int32_t token, int position, - std::vector & logits) { + std::vector & logits, + MoeHybridStreamEngine * stream_engine) { if (!backend || !w.ctx || !cache.ctx || position < 0 || position >= cache.max_ctx || position != cache.cur_pos || token < 0 || token >= w.n_vocab) { set_last_error("Kimi-K3 step: invalid backend, cache position, or token"); return false; } + if (w.routed_experts_streamed) { + if (!stream_engine || !stream_engine->is_bound()) { + set_last_error( + "Kimi-K3 step: file-backed experts require a bound stream engine"); + return false; + } + return streamed_kimi_k3_step( + backend, w, cache, token, position, + logits, *stream_engine); + } ggml_init_params params{}; params.mem_size = 64ull * 1024ull * 1024ull; diff --git a/server/src/kimi_k3/kimi_k3_internal.h b/server/src/kimi_k3/kimi_k3_internal.h index a1296bade..5b82a983e 100644 --- a/server/src/kimi_k3/kimi_k3_internal.h +++ b/server/src/kimi_k3/kimi_k3_internal.h @@ -11,6 +11,8 @@ #pragma once +#include "common/moe_hybrid_storage.h" + #include "ggml.h" #include "ggml-backend.h" @@ -21,6 +23,8 @@ namespace dflash::common { +class MoeHybridStreamEngine; + struct KimiK3Layer { bool recurrent = false; @@ -77,9 +81,21 @@ struct KimiK3Layer { }; struct KimiK3Weights { + // ctx/buf alias the first entries for compatibility with the original + // single-file loader. Split GGUFs retain one metadata context and optional + // resident backend buffer per shard. ggml_context * ctx = nullptr; ggml_backend_t backend = nullptr; ggml_backend_buffer_t buf = nullptr; + std::vector contexts; + std::vector buffers; + std::vector shard_paths; + + // The routed stacks may remain file-backed. Regions use MoE-layer-local + // indices [0, n_layer-n_dense_lead), independent of model layer numbers. + std::vector streamed_layer_regions; + size_t max_streamed_expert_bytes = 0; + bool routed_experts_streamed = false; ggml_tensor * tok_embd = nullptr; ggml_tensor * output_norm = nullptr; @@ -135,7 +151,8 @@ struct KimiK3Cache { bool load_kimi_k3_gguf(const std::string & path, ggml_backend_t backend, - KimiK3Weights & out); + KimiK3Weights & out, + bool stream_routed_experts = false); void free_kimi_k3_weights(KimiK3Weights & w); bool create_kimi_k3_cache(ggml_backend_t backend, @@ -154,6 +171,7 @@ bool kimi_k3_step(ggml_backend_t backend, KimiK3Cache & cache, int32_t token, int position, - std::vector & logits); + std::vector & logits, + MoeHybridStreamEngine * stream_engine = nullptr); } // namespace dflash::common diff --git a/server/src/kimi_k3/kimi_k3_loader.cpp b/server/src/kimi_k3/kimi_k3_loader.cpp index 9aa8a92f3..2919e9293 100644 --- a/server/src/kimi_k3/kimi_k3_loader.cpp +++ b/server/src/kimi_k3/kimi_k3_loader.cpp @@ -8,7 +8,11 @@ #include #include #include +#include +#include +#include #include +#include #include namespace dflash::common { @@ -26,7 +30,16 @@ uint32_t get_u32_or(const gguf_context * g, const char * key, uint32_t fallback) return fallback; } const gguf_type type = gguf_get_kv_type(g, id); + if (type == GGUF_TYPE_UINT8) return gguf_get_val_u8(g, id); + if (type == GGUF_TYPE_UINT16) return gguf_get_val_u16(g, id); if (type == GGUF_TYPE_UINT32) return gguf_get_val_u32(g, id); + if (type == GGUF_TYPE_UINT64) { + const uint64_t value = gguf_get_val_u64(g, id); + return value <= std::numeric_limits::max() + ? static_cast(value) : fallback; + } + if (type == GGUF_TYPE_INT8) return static_cast(gguf_get_val_i8(g, id)); + if (type == GGUF_TYPE_INT16) return static_cast(gguf_get_val_i16(g, id)); if (type == GGUF_TYPE_INT32) return static_cast(gguf_get_val_i32(g, id)); return fallback; } @@ -68,42 +81,195 @@ bool tensor_shape_is(const ggml_tensor * t, return t && t->ne[0] == ne0 && t->ne[1] == ne1 && t->ne[2] == ne2; } +bool is_routed_expert_tensor(const std::string & name) { + return name.find(".ffn_gate_exps.weight") != std::string::npos || + name.find(".ffn_up_exps.weight") != std::string::npos || + name.find(".ffn_down_exps.weight") != std::string::npos || + name.find(".ffn_gate_up_exps.weight") != std::string::npos; +} + +size_t align_up(size_t value, size_t alignment) { + if (alignment == 0) return value; + const size_t remainder = value % alignment; + return remainder == 0 ? value : value + (alignment - remainder); +} + +bool discover_split_paths(const std::string & supplied, + uint32_t split_count, + std::vector & out, + std::string & error) { + out.clear(); + const size_t of = supplied.rfind("-of-"); + if (split_count <= 1 && of == std::string::npos) { + out.push_back(supplied); + return true; + } + + if (of == std::string::npos) { + error = "split.count is greater than one but the GGUF filename has no -NNNNN-of-NNNNN suffix"; + return false; + } + const size_t index_dash = supplied.rfind('-', of - 1); + if (index_dash == std::string::npos || index_dash + 1 >= of) { + error = "cannot locate the split index in the GGUF filename"; + return false; + } + size_t total_end = of + 4; + while (total_end < supplied.size() && + supplied[total_end] >= '0' && supplied[total_end] <= '9') { + ++total_end; + } + const std::string index_text = + supplied.substr(index_dash + 1, of - index_dash - 1); + const std::string total_text = + supplied.substr(of + 4, total_end - (of + 4)); + if (index_text.empty() || total_text.empty() || + index_text.find_first_not_of("0123456789") != std::string::npos || + total_text.find_first_not_of("0123456789") != std::string::npos) { + error = "invalid split index/count in the GGUF filename"; + return false; + } + uint64_t filename_total = 0; + try { + filename_total = std::stoull(total_text); + } catch (...) { + error = "GGUF filename split count is not an integer"; + return false; + } + if (split_count != 0 && filename_total != split_count) { + error = "GGUF split.count disagrees with the filename"; + return false; + } + if (filename_total == 0 || + filename_total > std::numeric_limits::max()) { + error = "GGUF filename has an invalid split count"; + return false; + } + split_count = static_cast(filename_total); + + const std::string prefix = supplied.substr(0, index_dash + 1); + const std::string suffix = supplied.substr(total_end); + out.reserve(split_count); + for (uint32_t split = 1; split <= split_count; ++split) { + std::ostringstream path; + path << prefix << std::setw(static_cast(index_text.size())) + << std::setfill('0') << split + << "-of-" << total_text << suffix; + out.push_back(path.str()); + } + return true; +} + +struct TensorSource { + ggml_tensor * tensor = nullptr; + uint32_t shard = 0; + size_t file_offset = 0; + size_t file_size = 0; +}; + } // namespace bool load_kimi_k3_gguf(const std::string & path, ggml_backend_t backend, - KimiK3Weights & out) { + KimiK3Weights & out, + bool stream_routed_experts) { free_kimi_k3_weights(out); - ggml_context * meta_ctx = nullptr; - gguf_init_params params{}; - params.no_alloc = true; - params.ctx = &meta_ctx; - gguf_context * gctx = gguf_init_from_file(path.c_str(), params); - if (!gctx || !meta_ctx) { + ggml_context * first_meta = nullptr; + gguf_init_params first_params{}; + first_params.no_alloc = true; + first_params.ctx = &first_meta; + gguf_context * first_gguf = + gguf_init_from_file(path.c_str(), first_params); + if (!first_gguf || !first_meta) { set_last_error("Kimi-K3: failed to parse GGUF: " + path); - if (gctx) gguf_free(gctx); - if (meta_ctx) ggml_free(meta_ctx); + if (first_gguf) gguf_free(first_gguf); + if (first_meta) ggml_free(first_meta); return false; } + // Only the first file of a standard split GGUF is required to carry + // global metadata. If the caller supplied another shard, infer the count + // from its canonical filename and reopen the set in numerical order. + const uint32_t split_count = + get_u32_or(first_gguf, "split.count", 0); + std::vector shard_paths; + std::string discovery_error; + if (!discover_split_paths(path, split_count, shard_paths, + discovery_error)) { + gguf_free(first_gguf); + ggml_free(first_meta); + set_last_error("Kimi-K3: " + discovery_error); + return false; + } + + std::vector shard_ggufs; + shard_ggufs.reserve(shard_paths.size()); + out.contexts.reserve(shard_paths.size()); + bool used_supplied = false; + for (size_t shard = 0; shard < shard_paths.size(); ++shard) { + if (shard_paths[shard] == path) { + shard_ggufs.push_back(first_gguf); + out.contexts.push_back(first_meta); + used_supplied = true; + continue; + } + ggml_context * meta = nullptr; + gguf_init_params params{}; + params.no_alloc = true; + params.ctx = &meta; + gguf_context * gguf = + gguf_init_from_file(shard_paths[shard].c_str(), params); + if (!gguf || !meta) { + if (gguf) gguf_free(gguf); + if (meta) ggml_free(meta); + for (gguf_context * opened : shard_ggufs) gguf_free(opened); + if (!used_supplied) { + gguf_free(first_gguf); + ggml_free(first_meta); + } + free_kimi_k3_weights(out); + set_last_error("Kimi-K3: failed to parse GGUF shard: " + + shard_paths[shard]); + return false; + } + shard_ggufs.push_back(gguf); + out.contexts.push_back(meta); + } + if (!used_supplied) { + gguf_free(first_gguf); + ggml_free(first_meta); + } + auto fail = [&](const std::string & message) { set_last_error("Kimi-K3: " + message); - if (out.buf) { - ggml_backend_buffer_free(out.buf); - out.buf = nullptr; - } - gguf_free(gctx); - ggml_free(meta_ctx); - out = KimiK3Weights{}; + for (gguf_context * gguf : shard_ggufs) gguf_free(gguf); + shard_ggufs.clear(); + free_kimi_k3_weights(out); return false; }; - const int64_t arch_id = gguf_find_key(gctx, "general.architecture"); - if (arch_id < 0 || std::strcmp(gguf_get_val_str(gctx, arch_id), "kimi-k3") != 0) { - return fail("general.architecture must be kimi-k3"); + gguf_context * gctx = nullptr; + for (size_t shard = 0; shard < shard_ggufs.size(); ++shard) { + const int64_t arch_id = + gguf_find_key(shard_ggufs[shard], "general.architecture"); + if (arch_id >= 0) { + if (std::strcmp( + gguf_get_val_str(shard_ggufs[shard], arch_id), + "kimi-k3") != 0) { + return fail("general.architecture must be kimi-k3"); + } + if (!gctx) gctx = shard_ggufs[shard]; + } + const uint32_t count = + get_u32_or(shard_ggufs[shard], "split.count", 0); + if (count != 0 && count != shard_paths.size()) { + return fail("split.count is inconsistent across GGUF shards"); + } + } + if (!gctx) { + return fail("no shard contains general.architecture metadata"); } - constexpr const char * A = "kimi-k3."; auto key = [&](const char * suffix) { return std::string(A) + suffix; }; auto u32 = [&](const char * suffix, uint32_t fallback = 0) { @@ -119,8 +285,10 @@ bool load_kimi_k3_gguf(const std::string & path, return get_bool_or(gctx, k.c_str(), fallback); }; - out.ctx = meta_ctx; + out.ctx = out.contexts.front(); out.backend = backend; + out.shard_paths = shard_paths; + out.routed_experts_streamed = stream_routed_experts; out.n_layer = static_cast(u32("block_count")); out.n_embd = static_cast(u32("embedding_length")); out.n_ff = static_cast(u32("feed_forward_length")); @@ -150,7 +318,31 @@ bool load_kimi_k3_gguf(const std::string & path, out.situ_linear_beta = f32("activation.situ_linear_beta", 25.0f); out.eos_token_id = static_cast(get_u32_or(gctx, "tokenizer.ggml.eos_token_id", 2)); - auto get = [&](const char * name) { return ggml_get_tensor(meta_ctx, name); }; + std::unordered_map tensors; + for (size_t shard = 0; shard < shard_ggufs.size(); ++shard) { + gguf_context * gguf = shard_ggufs[shard]; + ggml_context * meta = out.contexts[shard]; + const size_t data_start = gguf_get_data_offset(gguf); + const int64_t count = gguf_get_n_tensors(gguf); + for (int64_t tid = 0; tid < count; ++tid) { + const char * name = gguf_get_tensor_name(gguf, tid); + if (!name || tensors.find(name) != tensors.end()) { + return fail(std::string("duplicate or unnamed tensor across shards: ") + + (name ? name : "")); + } + TensorSource source; + source.tensor = ggml_get_tensor(meta, name); + source.shard = static_cast(shard); + source.file_offset = + data_start + gguf_get_tensor_offset(gguf, tid); + source.file_size = gguf_get_tensor_size(gguf, tid); + tensors.emplace(name, source); + } + } + auto get = [&](const char * name) -> ggml_tensor * { + const auto found = tensors.find(name); + return found == tensors.end() ? nullptr : found->second.tensor; + }; out.tok_embd = get("token_embd.weight"); out.output_norm = get("output_norm.weight"); out.output = get("output.weight"); @@ -276,36 +468,136 @@ bool load_kimi_k3_gguf(const std::string & path, } } - out.buf = ggml_backend_alloc_ctx_tensors(meta_ctx, backend); - if (!out.buf) return fail("unable to allocate resident tensor buffer"); + out.streamed_layer_regions.clear(); + out.max_streamed_expert_bytes = 0; + for (int il = out.n_dense_lead; il < out.n_layer; ++il) { + char gate_name[160], up_name[160], down_name[160]; + std::snprintf(gate_name, sizeof(gate_name), + "blk.%d.ffn_gate_exps.weight", il); + std::snprintf(up_name, sizeof(up_name), + "blk.%d.ffn_up_exps.weight", il); + std::snprintf(down_name, sizeof(down_name), + "blk.%d.ffn_down_exps.weight", il); + const TensorSource & gate = tensors.at(gate_name); + const TensorSource & up = tensors.at(up_name); + const TensorSource & down = tensors.at(down_name); + if (gate.file_size % static_cast(out.n_expert) != 0 || + up.file_size % static_cast(out.n_expert) != 0 || + down.file_size % static_cast(out.n_expert) != 0) { + return fail("routed expert tensor size is not divisible by expert_count"); + } + LayerExpertRegions regions; + regions.expert_bytes_gate = + gate.file_size / static_cast(out.n_expert); + regions.expert_bytes_up = + up.file_size / static_cast(out.n_expert); + regions.expert_bytes_down = + down.file_size / static_cast(out.n_expert); + regions.gate_exps = { + gate.file_offset, gate.file_size, gate.shard}; + regions.up_exps = { + up.file_offset, up.file_size, up.shard}; + regions.down_exps = { + down.file_offset, down.file_size, down.shard}; + const size_t expert_bytes = + regions.expert_bytes_gate + regions.expert_bytes_up + + regions.expert_bytes_down; + out.max_streamed_expert_bytes = + std::max(out.max_streamed_expert_bytes, expert_bytes); + out.streamed_layer_regions.push_back(regions); + } - GgufMmap mmap; - std::string mmap_error; - if (!mmap.open(path, mmap_error)) return fail(mmap_error); - const auto * base = static_cast(mmap.data()); - const size_t file_size = mmap.size(); - const size_t data_start = gguf_get_data_offset(gctx); + struct ResidentAlloc { + ggml_tensor * tensor = nullptr; + size_t file_offset = 0; + size_t file_size = 0; + size_t buffer_offset = 0; + }; + ggml_backend_buffer_type_t buft = + ggml_backend_get_default_buffer_type(backend); + const size_t alignment = ggml_backend_buft_get_alignment(buft); size_t copied = 0; - for (int64_t tid = 0; tid < gguf_get_n_tensors(gctx); ++tid) { - const char * tensor_name = gguf_get_tensor_name(gctx, tid); - ggml_tensor * tensor = ggml_get_tensor(meta_ctx, tensor_name); - if (!tensor) continue; - const size_t offset = gguf_get_tensor_offset(gctx, tid); - const size_t bytes = gguf_get_tensor_size(gctx, tid); - if (!gguf_tensor_in_file(data_start, offset, bytes, file_size)) { - return fail(gguf_bounds_error("Kimi-K3 GGUF", tensor_name, - ggml_type_name(gguf_get_tensor_type(gctx, tid)), data_start, - offset, bytes, file_size)); + size_t skipped = 0; + for (size_t shard = 0; shard < shard_ggufs.size(); ++shard) { + gguf_context * gguf = shard_ggufs[shard]; + ggml_context * meta = out.contexts[shard]; + const size_t data_start = gguf_get_data_offset(gguf); + std::vector allocs; + size_t allocation_bytes = 0; + for (int64_t tid = 0; tid < gguf_get_n_tensors(gguf); ++tid) { + const char * tensor_name = gguf_get_tensor_name(gguf, tid); + ggml_tensor * tensor = ggml_get_tensor(meta, tensor_name); + const size_t bytes = gguf_get_tensor_size(gguf, tid); + if (!tensor) continue; + if (stream_routed_experts && + is_routed_expert_tensor(tensor_name)) { + skipped += bytes; + continue; + } + allocation_bytes = align_up(allocation_bytes, alignment); + ResidentAlloc allocation; + allocation.tensor = tensor; + allocation.file_offset = + data_start + gguf_get_tensor_offset(gguf, tid); + allocation.file_size = bytes; + allocation.buffer_offset = allocation_bytes; + allocation_bytes += + ggml_backend_buft_get_alloc_size(buft, tensor); + allocs.push_back(allocation); + } + if (allocs.empty()) continue; + + ggml_backend_buffer_t buffer = + ggml_backend_alloc_buffer(backend, allocation_bytes); + if (!buffer) { + return fail("unable to allocate resident tensor buffer for shard " + + std::to_string(shard + 1)); + } + ggml_backend_buffer_set_usage( + buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + out.buffers.push_back(buffer); + char * buffer_base = + static_cast(ggml_backend_buffer_get_base(buffer)); + for (const ResidentAlloc & allocation : allocs) { + if (ggml_backend_tensor_alloc( + buffer, allocation.tensor, + buffer_base + allocation.buffer_offset) != + GGML_STATUS_SUCCESS) { + return fail("unable to bind a resident tensor allocation"); + } + } + + GgufMmap mmap; + std::string mmap_error; + if (!mmap.open(shard_paths[shard], mmap_error)) { + return fail(mmap_error); + } + const auto * base = + static_cast(mmap.data()); + for (const ResidentAlloc & allocation : allocs) { + if (allocation.file_offset > mmap.size() || + allocation.file_size > + mmap.size() - allocation.file_offset) { + return fail("resident tensor range is outside GGUF shard " + + std::to_string(shard + 1)); + } + ggml_backend_tensor_set( + allocation.tensor, base + allocation.file_offset, + 0, allocation.file_size); + copied += allocation.file_size; } - ggml_backend_tensor_set(tensor, base + data_start + offset, 0, bytes); - copied += bytes; } + out.buf = out.buffers.empty() ? nullptr : out.buffers.front(); - gguf_free(gctx); + for (gguf_context * gguf : shard_ggufs) gguf_free(gguf); + shard_ggufs.clear(); std::fprintf(stderr, - "[kimi-k3] loaded %.2f GiB: layers=%d (KDA=%zu MLA=%zu) hidden=%d " + "[kimi-k3] loaded resident=%.2f GiB file-backed-experts=%.2f GiB " + "shards=%zu layers=%d (KDA=%zu MLA=%zu) hidden=%d " "experts=%d top=%d latent=%d vocab=%d\n", static_cast(copied) / (1024.0 * 1024.0 * 1024.0), + static_cast(skipped) / (1024.0 * 1024.0 * 1024.0), + out.shard_paths.size(), out.n_layer, static_cast(std::count_if(out.layers.begin(), out.layers.end(), [](const KimiK3Layer & l) { return l.recurrent; })), @@ -318,8 +610,12 @@ bool load_kimi_k3_gguf(const std::string & path, } void free_kimi_k3_weights(KimiK3Weights & w) { - if (w.buf) ggml_backend_buffer_free(w.buf); - if (w.ctx) ggml_free(w.ctx); + for (ggml_backend_buffer_t buffer : w.buffers) { + if (buffer) ggml_backend_buffer_free(buffer); + } + for (ggml_context * context : w.contexts) { + if (context) ggml_free(context); + } w = KimiK3Weights{}; } diff --git a/server/test/smoke_kimi_k3_forward.cpp b/server/test/smoke_kimi_k3_forward.cpp index 9c6b42c8f..cb71f1cf2 100644 --- a/server/test/smoke_kimi_k3_forward.cpp +++ b/server/test/smoke_kimi_k3_forward.cpp @@ -10,7 +10,8 @@ using namespace dflash::common; int main(int argc, char ** argv) { if (argc < 2) { std::fprintf(stderr, - "usage: %s [gpu=0] [n_gen=16] [prompt]\n", + "usage: %s [gpu=0] [n_gen=16] [prompt] " + "[stream_experts=1]\n", argv[0]); return 2; } @@ -36,6 +37,8 @@ int main(int argc, char ** argv) { config.model_path = model; config.device.gpu = gpu; config.device.max_ctx = 4096; + config.stream_routed_experts = + argc <= 5 || std::atoi(argv[5]) != 0; KimiK3Backend backend(config); if (!backend.init()) return 1; From c42a254489ad1c0ceec4bbd5e39ae1edc35cb190 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:34:56 +0200 Subject: [PATCH 09/20] feat(moe): add reusable expert GPU ownership --- server/CMakeLists.txt | 1 + server/docs/ENVIRONMENT.md | 1 + server/docs/KIMI_K3_HETERO.md | 29 ++++--- server/docs/MOE_NVME_STREAMING.md | 27 +++++- server/src/common/moe_hybrid_placement.cpp | 42 ++++++++++ server/src/common/moe_hybrid_placement.h | 19 +++++ server/src/deepseek4/deepseek4_backend.cpp | 28 +++++-- server/src/kimi_k3/kimi_k3_backend.cpp | 65 +++++++++++--- server/src/kimi_k3/kimi_k3_backend.h | 7 ++ server/test/smoke_kimi_k3_forward.cpp | 3 +- .../test/test_moe_expert_owner_placement.cpp | 84 +++++++++++++++++++ .../test/test_qwen35moe_expert_placement.cpp | 1 - variables.md | 1 + 13 files changed, 270 insertions(+), 38 deletions(-) create mode 100644 server/test/test_moe_expert_owner_placement.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index eaa0f1659..9a8f94cf4 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1183,6 +1183,7 @@ if(DFLASH27B_TESTS) test/test_kvflash_placement.cpp test/test_kvflash_pool_sizing.cpp test/test_kvflash_qk.cpp + test/test_moe_expert_owner_placement.cpp test/test_qwen35moe_routing_stats.cpp test/test_qwen35moe_expert_placement.cpp test/test_qwen35moe_swap_manager.cpp diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index 5cdafe579..e0c1e1ba8 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -28,6 +28,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_MMID_GROUPED_DEVICE` | -1 | Optional zero-based device restriction; unset/-1 applies to every eligible device. | | `DFLASH_DS4_MOE_TP` / `DFLASH_DS4_MOE_TP_INPROC` | unset | BURN-IN: enable DeepSeek4 route-owner expert parallelism in one process. | | `DFLASH_DS4_MOE_TP_GPU` | auto | HIP device for the cold DeepSeek4 expert owner. | +| `DFLASH_MOE_TP_GPU` | primary device | Model-neutral routed-expert owner GPU; used by Kimi streaming and preferred over legacy DS4/IPC GPU variables. | | `DFLASH_MOE_NVME_COLD_TIER` | auto | BURN-IN: DeepSeek4 cold-capacity policy (`auto`, `on`, `off`) for dual-device and Strix-only execution. | | `DFLASH_MOE_NVME_*` | tuned defaults | BURN-IN: bounded MoE SSD scheduler/backend controls; see `MOE_NVME_STREAMING.md`. | | `GGML_CUDA_BATCH_PEER_COPIES` | unset | BURN-IN: publish ordered HIP peer copies with one cross-device dependency per source/destination pair. | diff --git a/server/docs/KIMI_K3_HETERO.md b/server/docs/KIMI_K3_HETERO.md index 989c4997e..978524981 100644 --- a/server/docs/KIMI_K3_HETERO.md +++ b/server/docs/KIMI_K3_HETERO.md @@ -29,6 +29,10 @@ small architecture adapter around it: are allocated. It reserves the larger of 2 GiB or 5% of device memory and is capped by the actual routed pool, so a tiny model cannot accidentally request a model-sized cache. +- Routed-expert ownership is independent of contiguous layer splitting. + `DFLASH_MOE_TP_GPU` can place the SSD slots, adaptive cache, and selected + expert graphs on a second HIP device while KDA/MLA, recurrent/KV state, + shared experts, and sampling remain on the primary device. - `bench_kimi_k3_hetero` runs Kimi's exact IQ1_S routed-expert geometry through the same model-neutral persistent evaluator used by production adapters: 896 experts, top-16, 92 MoE layers, @@ -45,10 +49,9 @@ The SSD text path is implemented, but two qualification boundaries remain: - The backend is correctness-first and token-sequential. Its per-layer graph boundaries are not yet fused/captured for full-model speed. -- Kimi currently places resident text tensors and streamed expert compute on - one selected GPU. Strix-only is supported directly. Splitting Kimi's dense - tensors onto R9700 while Strix owns streamed experts is a later performance - adapter, not a requirement for SSD capacity correctness. +- The heterogeneous path currently crosses a host-visible, activation-sized + boundary at every routed layer. It is correct on two devices but does not + yet use PR-505's device-resident owner fork/join and overlap. The vision encoder is out of scope for this text-only path. @@ -126,18 +129,19 @@ a cache only helps in proportion to its share of the 495 GiB routed pool. The machine has about 125.08 GiB of system/UMA memory plus 31.86 GiB on the R9700, or 156.94 GiB of unique physical weight capacity before runtime -reserves. The implemented capacity-safe starting point is: +reserves. The implemented capacity-safe Strix-only starting point is: 1. Strix/system memory: all non-routed text weights, shared experts, latent projections, recurrent/KV state, workspace, and the routed-expert cache. 2. NVMe: all routed expert stacks, with actual route misses read directly into pinned slots and evaluated on Strix. -This works unchanged on a Strix-only machine. On a full Lucebox the R9700 is -currently unused by the Kimi adapter; moving dense KDA/MLA work there is the -next throughput optimization. The common SSD runtime already supports a -different compute owner, but Kimi still needs device-aware resident allocation -and cache ownership before that placement is correct end to end. +This works unchanged on a Strix-only machine. A full Lucebox can additionally +assign the R9700 as the routed-expert owner with `DFLASH_MOE_TP_GPU`, moving +the adaptive expert cache and selected expert compute off Strix. That uses the +second device's memory, but the measured R9700 cold-storage path is slower, so +it is a capacity option rather than the default speed profile. The inverse +placement is also supported when a model's non-routed tensors fit on R9700. After approximately 57.94 GiB of non-routed weights plus OS, workspace, and a moderate context reserve, roughly 70-85 GiB may remain for routed experts. @@ -164,9 +168,8 @@ router. Full-scale qualification requires: chosen context budget. 3. Compare end-to-end output and token rate against ordinary llama.cpp CPU/GPU offload. -4. Only after that baseline, place dense attention on R9700 and the - latent/shared/routed MoE path on Strix, retaining activation-sized transfers - at the boundary. +4. Only after that baseline, replace the correctness-first host boundary with + a device-resident owner fork/join and overlap dense work with expert service. The full 594 GB model cannot currently be staged on the qualification box, which currently has about 513 GB free. It needs at least about 650 GB of safe diff --git a/server/docs/MOE_NVME_STREAMING.md b/server/docs/MOE_NVME_STREAMING.md index 52e17f250..1f3111c3c 100644 --- a/server/docs/MOE_NVME_STREAMING.md +++ b/server/docs/MOE_NVME_STREAMING.md @@ -139,8 +139,25 @@ the complete routed pool. `DFLASH_MOE_NVME_DEVICE_CACHE_MB` remains an explicit override, also capped by the routed pool. Kimi's native router remains authoritative. The current text backend is -correctness-first and sequential; multi-device dense placement, captured -per-layer graphs, and the vision tower are separate optimizations. +correctness-first and sequential; captured per-layer graphs and the vision +tower are separate optimizations. + +On a two-GPU Lucebox, the primary device still owns KDA/MLA and model state; +the optional expert device owns the SSD slots, adaptive cache, and selected +expert graphs: + +```bash +export DFLASH_MOE_TP_GPU= + +./build-hip-dual/dflash_server \ + /path/to/Kimi-K3-UD-IQ1_S-00001-of-00014.gguf \ + --target-device hip: --max-ctx 8192 +``` + +This is functional expert ownership, not a contiguous layer split, so do not +use `--target-devices`. Leaving `DFLASH_MOE_TP_GPU` unset preserves the +Strix-only path. The current Kimi boundary uses activation-sized host staging; +device-resident peer fork/join is a later throughput optimization. ## Tuning and diagnostics @@ -160,6 +177,7 @@ not improve the qualified P310 drive and consumes extra pinned/system memory. | `DFLASH_MOE_NVME_DEVICE_CACHE_MB` | automatic/`0` | Override adaptive device expert-cache memory; `0` leaves only pipeline slots | | `DFLASH_MOE_NVME_GRAPH_CACHE` | `8` | Persistent expert-graph variants retained per stream engine; `0` is a diagnostic no-cache mode | | `DFLASH_MOE_NVME_REFERENCE_EVAL` | unset | Diagnostic only: `1` restores the allocation-heavy reference evaluator for numerical/performance A/B | +| `DFLASH_MOE_TP_GPU` | primary GPU | Optional second GPU that owns streamed expert cache and compute | Shutdown telemetry reports logical and physical bytes, measured read service rate, cache hits, demand wait, de-duplication, dropped speculation, errors, and @@ -172,8 +190,9 @@ targets `test_moe_nvme_scheduler`, `bench_moe_nvme_io`, and SSD-to-GPU path. Benchmarks are read-only. The external `smoke_kimi_k3_forward` target accepts -`[stream_experts=0|1]`. This is an A/B oracle for small Kimi fixtures; production -Kimi uses `1`. +`[stream_experts=0|1] [expert_gpu=-1]`. This is an A/B and placement oracle for +small Kimi fixtures; production Kimi uses streaming and `-1` resolves the +environment/default owner. ## Qualification result (2026-07-30) diff --git a/server/src/common/moe_hybrid_placement.cpp b/server/src/common/moe_hybrid_placement.cpp index 49bcd5ea1..73130b6f9 100644 --- a/server/src/common/moe_hybrid_placement.cpp +++ b/server/src/common/moe_hybrid_placement.cpp @@ -4,11 +4,53 @@ #include #include +#include +#include #include +#include #include namespace dflash::common { +bool resolve_moe_expert_owner_placement( + int primary_gpu, + int requested_expert_gpu, + MoeExpertOwnerPlacement & out, + std::string * err) { + if (primary_gpu < 0 || requested_expert_gpu < -1) { + if (err) *err = "MoE owner GPU indices must be non-negative"; + return false; + } + + int expert_gpu = requested_expert_gpu; + if (expert_gpu < 0) { + const char * raw = std::getenv("DFLASH_MOE_TP_GPU"); + if (!raw || !*raw) raw = std::getenv("DFLASH_DS4_MOE_TP_GPU"); + if (!raw || !*raw) { + raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU"); + } + if (!raw || !*raw) { + expert_gpu = primary_gpu; + } else { + errno = 0; + char * end = nullptr; + const long parsed = std::strtol(raw, &end, 10); + if (errno != 0 || end == raw || *end != '\0' || parsed < 0 || + parsed > std::numeric_limits::max()) { + if (err) { + *err = std::string("invalid routed-expert GPU: ") + raw; + } + return false; + } + expert_gpu = static_cast(parsed); + } + } + + out.primary_gpu = primary_gpu; + out.expert_gpu = expert_gpu; + return true; +} + bool MoeHybridPlacement::matches(int n_layer_, int n_expert_, int n_expert_used_) const { return n_layer == n_layer_ && n_expert == n_expert_ && diff --git a/server/src/common/moe_hybrid_placement.h b/server/src/common/moe_hybrid_placement.h index a522a4825..902c91321 100644 --- a/server/src/common/moe_hybrid_placement.h +++ b/server/src/common/moe_hybrid_placement.h @@ -14,6 +14,25 @@ namespace dflash::common { struct MoeHybridRoutingStats; // forward decl +// Functional MoE placement: the model's dense/recurrent graph remains on the +// primary owner while selected routed experts may execute on a second GPU. +// This is deliberately separate from contiguous layer splitting. +struct MoeExpertOwnerPlacement { + int primary_gpu = 0; + int expert_gpu = 0; + + bool heterogeneous() const { return primary_gpu != expert_gpu; } +}; + +// Resolve a reusable routed-expert owner. requested_expert_gpu >= 0 is an +// explicit programmatic choice. -1 reads DFLASH_MOE_TP_GPU and then legacy +// DeepSeek/IPC spellings; if none is set, experts stay on primary_gpu. +bool resolve_moe_expert_owner_placement( + int primary_gpu, + int requested_expert_gpu, + MoeExpertOwnerPlacement & out, + std::string * err = nullptr); + inline uint64_t moe_hybrid_core_bytes_from_memory(const char * log_prefix, size_t gpu_free, size_t gpu_total) { diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 0c00ff228..ea4d831c1 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -117,12 +117,25 @@ static bool ds4_inprocess_moe_tp_enabled() { } static int ds4_moe_tp_gpu(int local_gpu) { - const char * raw = std::getenv("DFLASH_DS4_MOE_TP_GPU"); - if (!raw || !*raw) { - raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU"); + const char * generic = std::getenv("DFLASH_MOE_TP_GPU"); + const char * legacy_ds4 = std::getenv("DFLASH_DS4_MOE_TP_GPU"); + const char * legacy_ipc = + std::getenv("DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU"); + const bool configured = + (generic && *generic) || (legacy_ds4 && *legacy_ds4) || + (legacy_ipc && *legacy_ipc); + const int default_peer = local_gpu == 0 ? 1 : 0; + MoeExpertOwnerPlacement owner; + std::string error; + if (!resolve_moe_expert_owner_placement( + local_gpu, configured ? -1 : default_peer, + owner, &error)) { + std::fprintf(stderr, + "[deepseek4-moe-tp] invalid expert owner: %s\n", + error.c_str()); + return -1; } - if (raw && *raw) return std::max(0, std::atoi(raw)); - return local_gpu == 0 ? 1 : 0; + return owner.expert_gpu; } static double gib(uint64_t bytes) { @@ -789,6 +802,8 @@ bool DeepSeek4Backend::init_moe_tensor_parallel() { } if (ds4_inprocess_moe_tp_enabled()) { + const int expert_gpu = ds4_moe_tp_gpu(cfg_.device.gpu); + if (expert_gpu < 0) return false; const bool resident_ready = moe_hybrid_->materialized_cold_experts && moe_hybrid_->cold_backend == expert_backend_; const bool streamed_ready = !moe_hybrid_->materialized_cold_experts && @@ -804,7 +819,7 @@ bool DeepSeek4Backend::init_moe_tensor_parallel() { "[deepseek4-moe-tp] enabled mode=%s local_gpu=%d " "expert_gpu=%d local_experts=%d remote_experts=%d\n", streamed_ready ? "in-process+ssd" : "in-process", - cfg_.device.gpu, ds4_moe_tp_gpu(cfg_.device.gpu), + cfg_.device.gpu, expert_gpu, moe_placement_.total_hot, w_.n_layer * w_.n_expert - moe_placement_.total_hot); return true; @@ -951,6 +966,7 @@ bool DeepSeek4Backend::init_hybrid_model() { nvme_mode != MoeNvmeColdTierMode::Disabled; if (inprocess_tp) { const int expert_gpu = ds4_moe_tp_gpu(cfg_.device.gpu); + if (expert_gpu < 0) return false; if (expert_gpu == cfg_.device.gpu) { std::fprintf(stderr, "[deepseek4-moe-tp] in-process expert GPU must differ from local GPU\n"); diff --git a/server/src/kimi_k3/kimi_k3_backend.cpp b/server/src/kimi_k3/kimi_k3_backend.cpp index 269862697..6b2ab1611 100644 --- a/server/src/kimi_k3/kimi_k3_backend.cpp +++ b/server/src/kimi_k3/kimi_k3_backend.cpp @@ -1,5 +1,6 @@ #include "kimi_k3_backend.h" +#include "common/moe_hybrid_placement.h" #include "common/sampler.h" #include "dflash27b.h" @@ -33,6 +34,14 @@ KimiK3Backend::~KimiK3Backend() { shutdown(); } +void KimiK3Backend::release_expert_backend() { + if (expert_backend_) { + ggml_backend_free(expert_backend_); + expert_backend_ = nullptr; + } + expert_gpu_ = -1; +} + bool KimiK3Backend::init_streaming() { if (!weights_.routed_experts_streamed || weights_.streamed_layer_regions.empty() || @@ -42,6 +51,35 @@ bool KimiK3Backend::init_streaming() { return false; } + MoeExpertOwnerPlacement owner; + std::string error; + if (!resolve_moe_expert_owner_placement( + cfg_.device.primary_gpu(), cfg_.expert_gpu, + owner, &error)) { + std::fprintf(stderr, + "[kimi-k3] invalid expert-owner placement: %s\n", + error.c_str()); + return false; + } + expert_gpu_ = owner.expert_gpu; + if (owner.heterogeneous()) { + expert_backend_ = ggml_backend_cuda_init(expert_gpu_); + if (!expert_backend_) { + std::fprintf(stderr, + "[kimi-k3] expert backend init failed for device %d\n", + expert_gpu_); + expert_gpu_ = -1; + return false; + } + } + ggml_backend_t stream_backend = + expert_backend_ ? expert_backend_ : backend_; + auto fail_streaming = [&]() { + stream_engine_.destroy(); + release_expert_backend(); + return false; + }; + MoeStreamConfig stream_config = MoeStreamConfig::from_env(); size_t routed_pool_bytes = 0; for (const LayerExpertRegions & regions : @@ -75,7 +113,7 @@ bool KimiK3Backend::init_streaming() { size_t free_bytes = 0; size_t total_bytes = 0; ggml_backend_cuda_get_device_memory( - cfg_.device.primary_gpu(), &free_bytes, &total_bytes); + expert_gpu_, &free_bytes, &total_bytes); const size_t gib = 1024ULL * 1024ULL * 1024ULL; const size_t reserve = std::max(2 * gib, total_bytes / 20); stream_config.device_cache_bytes = @@ -93,14 +131,13 @@ bool KimiK3Backend::init_streaming() { stream_config.device_cache_bytes = std::min(stream_config.device_cache_bytes, routed_pool_bytes); } - std::string error; if (!stream_engine_.init( - backend_, weights_.max_streamed_expert_bytes, + stream_backend, weights_.max_streamed_expert_bytes, stream_config, &error)) { std::fprintf(stderr, "[kimi-k3] stream engine initialization failed: %s\n", error.c_str()); - return false; + return fail_streaming(); } std::vector descriptors; @@ -124,8 +161,7 @@ bool KimiK3Backend::init_streaming() { ::close(opened); #endif } - stream_engine_.destroy(); - return false; + return fail_streaming(); } uint64_t shard_bytes = 0; #if defined(_WIN32) @@ -156,8 +192,7 @@ bool KimiK3Backend::init_streaming() { ::close(opened); #endif } - stream_engine_.destroy(); - return false; + return fail_streaming(); } descriptors.push_back(fd); sources.push_back({ @@ -176,15 +211,15 @@ bool KimiK3Backend::init_streaming() { std::fprintf(stderr, "[kimi-k3] stream source binding failed: %s\n", error.c_str()); - stream_engine_.destroy(); - return false; + return fail_streaming(); } std::fprintf(stderr, "[kimi-k3] routed experts file-backed: shards=%zu layers=%zu " - "io=%s cache=%.2f GiB\n", + "io=%s primary_gpu=%d expert_gpu=%d cache=%.2f GiB\n", weights_.shard_paths.size(), weights_.streamed_layer_regions.size(), stream_engine_.io_backend_name(), + cfg_.device.primary_gpu(), expert_gpu_, static_cast(stream_engine_.device_cache_bytes()) / (1024.0 * 1024.0 * 1024.0)); return true; @@ -217,9 +252,11 @@ bool KimiK3Backend::init() { if (weights_.routed_experts_streamed && !init_streaming()) return false; std::fprintf(stderr, "[kimi-k3] native backend ready on device %d (max_ctx=%d, " - "experts=%s, correctness-first sequential prefill)\n", + "experts=%s:%d, correctness-first sequential prefill)\n", cfg_.device.primary_gpu(), max_ctx, - weights_.routed_experts_streamed ? "nvme" : "resident"); + weights_.routed_experts_streamed ? "nvme" : "resident", + weights_.routed_experts_streamed ? expert_gpu_ + : cfg_.device.primary_gpu()); std::fflush(stderr); return true; } @@ -236,6 +273,7 @@ bool KimiK3Backend::park(ParkTarget target) { if (!park_target_includes_target_model(target)) return false; if (!parked_) { stream_engine_.destroy(); + release_expert_backend(); free_kimi_k3_weights(weights_); parked_ = true; } @@ -391,6 +429,7 @@ bool KimiK3Backend::handle_compress(const std::string & line, void KimiK3Backend::shutdown() { stream_engine_.destroy(); + release_expert_backend(); free_kimi_k3_cache(cache_); free_kimi_k3_weights(weights_); if (backend_) { diff --git a/server/src/kimi_k3/kimi_k3_backend.h b/server/src/kimi_k3/kimi_k3_backend.h index 33b57b000..6f79cd493 100644 --- a/server/src/kimi_k3/kimi_k3_backend.h +++ b/server/src/kimi_k3/kimi_k3_backend.h @@ -14,6 +14,10 @@ struct KimiK3BackendConfig { const char * model_path = nullptr; DevicePlacement device; int stream_fd = -1; + // -1 resolves DFLASH_MOE_TP_GPU and otherwise keeps experts on the + // primary GPU. A different GPU owns streamed weights/cache/compute while + // dense KDA/MLA, recurrent state, and sampling remain primary-owned. + int expert_gpu = -1; // Production Kimi uses file-backed routed experts. The resident mode is // retained as a deterministic oracle for small architecture fixtures. bool stream_routed_experts = true; @@ -49,6 +53,7 @@ class KimiK3Backend final : public ModelBackend { private: bool init_streaming(); + void release_expert_backend(); int32_t choose_token(const std::vector & logits, const SamplerCfg & sampler, @@ -56,6 +61,8 @@ class KimiK3Backend final : public ModelBackend { KimiK3BackendConfig cfg_; ggml_backend_t backend_ = nullptr; + ggml_backend_t expert_backend_ = nullptr; + int expert_gpu_ = -1; KimiK3Weights weights_; KimiK3Cache cache_; MoeHybridStreamEngine stream_engine_; diff --git a/server/test/smoke_kimi_k3_forward.cpp b/server/test/smoke_kimi_k3_forward.cpp index cb71f1cf2..e8d258768 100644 --- a/server/test/smoke_kimi_k3_forward.cpp +++ b/server/test/smoke_kimi_k3_forward.cpp @@ -11,7 +11,7 @@ int main(int argc, char ** argv) { if (argc < 2) { std::fprintf(stderr, "usage: %s [gpu=0] [n_gen=16] [prompt] " - "[stream_experts=1]\n", + "[stream_experts=1] [expert_gpu=-1]\n", argv[0]); return 2; } @@ -39,6 +39,7 @@ int main(int argc, char ** argv) { config.device.max_ctx = 4096; config.stream_routed_experts = argc <= 5 || std::atoi(argv[5]) != 0; + config.expert_gpu = argc > 6 ? std::atoi(argv[6]) : -1; KimiK3Backend backend(config); if (!backend.init()) return 1; diff --git a/server/test/test_moe_expert_owner_placement.cpp b/server/test/test_moe_expert_owner_placement.cpp new file mode 100644 index 000000000..ed6d53fde --- /dev/null +++ b/server/test/test_moe_expert_owner_placement.cpp @@ -0,0 +1,84 @@ +#include "CppUnitTestFramework.hpp" +#include "../src/common/moe_hybrid_placement.h" +#include "../src/common/platform_env.h" + +#include +#include + +using namespace dflash::common; + +namespace { + +struct MoeExpertOwnerPlacementFixture {}; + +class ScopedEnvironment { +public: + explicit ScopedEnvironment(const char * name) : name_(name) { + const char * value = std::getenv(name); + if (value) { + existed_ = true; + value_ = value; + } + } + + ~ScopedEnvironment() { + if (existed_) { + set_environment_variable(name_.c_str(), value_.c_str(), true); + } else { + unset_environment_variable(name_.c_str()); + } + } + +private: + std::string name_; + std::string value_; + bool existed_ = false; +}; + +} // namespace + +TEST_CASE(MoeExpertOwnerPlacementFixture, resolves_programmatic_and_environment_owners) { + ScopedEnvironment generic_gpu("DFLASH_MOE_TP_GPU"); + ScopedEnvironment ds4_gpu("DFLASH_DS4_MOE_TP_GPU"); + ScopedEnvironment ipc_gpu("DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU"); + REQUIRE(unset_environment_variable("DFLASH_MOE_TP_GPU") == 0); + REQUIRE(unset_environment_variable("DFLASH_DS4_MOE_TP_GPU") == 0); + REQUIRE(unset_environment_variable( + "DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU") == 0); + + MoeExpertOwnerPlacement owner; + std::string error; + REQUIRE(resolve_moe_expert_owner_placement(0, -1, owner, &error)); + REQUIRE(owner.primary_gpu == 0); + REQUIRE(owner.expert_gpu == 0); + REQUIRE(!owner.heterogeneous()); + + REQUIRE(resolve_moe_expert_owner_placement(0, 1, owner, &error)); + REQUIRE(owner.expert_gpu == 1); + REQUIRE(owner.heterogeneous()); + REQUIRE(!resolve_moe_expert_owner_placement(-1, 0, owner, &error)); + REQUIRE(!resolve_moe_expert_owner_placement(0, -2, owner, &error)); + + REQUIRE(set_environment_variable("DFLASH_MOE_TP_GPU", "2", true) == 0); + REQUIRE(set_environment_variable("DFLASH_DS4_MOE_TP_GPU", "3", true) == 0); + REQUIRE(set_environment_variable( + "DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU", "4", true) == 0); + REQUIRE(resolve_moe_expert_owner_placement(0, -1, owner, &error)); + REQUIRE(owner.expert_gpu == 2); + + REQUIRE(unset_environment_variable("DFLASH_MOE_TP_GPU") == 0); + REQUIRE(resolve_moe_expert_owner_placement(0, -1, owner, &error)); + REQUIRE(owner.expert_gpu == 3); + REQUIRE(unset_environment_variable("DFLASH_DS4_MOE_TP_GPU") == 0); + REQUIRE(resolve_moe_expert_owner_placement(0, -1, owner, &error)); + REQUIRE(owner.expert_gpu == 4); + + REQUIRE(set_environment_variable( + "DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU", + "not-a-device", true) == 0); + REQUIRE(!resolve_moe_expert_owner_placement(0, -1, owner, &error)); + + // An explicit API choice must not be overridden by process environment. + REQUIRE(resolve_moe_expert_owner_placement(0, 1, owner, &error)); + REQUIRE(owner.expert_gpu == 1); +} diff --git a/server/test/test_qwen35moe_expert_placement.cpp b/server/test/test_qwen35moe_expert_placement.cpp index 459925866..734823309 100644 --- a/server/test/test_qwen35moe_expert_placement.cpp +++ b/server/test/test_qwen35moe_expert_placement.cpp @@ -3,7 +3,6 @@ #include "../src/common/moe_hybrid_routing_stats.h" #include -#include #include #include diff --git a/variables.md b/variables.md index fb06395b6..822097ff7 100644 --- a/variables.md +++ b/variables.md @@ -217,6 +217,7 @@ Untagged variables are operational tuning knobs. | Variable | Purpose | |---|---| +| `DFLASH_MOE_TP_GPU` | Model-neutral HIP device that owns routed-expert cache and compute; Kimi uses the primary device when unset. | | `DFLASH_MOE_HYBRID_PREFILL_EAGER` / `DFLASH_MOE_PREFILL_TRACE` | Model-neutral heterogeneous prefill policy and tracing. The legacy `DFLASH_DS4_*` spellings remain aliases. | | `DFLASH_MOE_TP_GROUPED_MMVQ` / `DFLASH_MOE_TP_FUSED_GATE_UP` | Model-neutral grouped and fused routed-FFN kernel qualification switches. The legacy `DFLASH_DS4_*` spellings remain aliases. | | `DFLASH_MOE_TP_COARSE_OWNER` / `DFLASH_MOE_TP_COARSE_OWNER_SPLIT` | Model-neutral owner-op lowering switches. The legacy `DFLASH_DS4_*` spellings remain aliases. | From 9796c050dd8500927351ddf8a56ca3a64e9b0100 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:30:21 +0200 Subject: [PATCH 10/20] feat(moe): overlap SSD experts across GPU owners --- server/CMakeLists.txt | 1 + server/docs/ENVIRONMENT.md | 8 +- server/docs/KIMI_K3_HETERO.md | 37 +- server/docs/MOE_NVME_STREAMING.md | 72 +++- server/src/common/moe_hybrid_stream.cpp | 378 ++++++++++++++++++ server/src/common/moe_hybrid_stream.h | 73 ++++ server/src/kimi_k3/kimi_k3_backend.cpp | 147 +++++-- server/src/kimi_k3/kimi_k3_backend.h | 9 +- server/src/kimi_k3/kimi_k3_graph.cpp | 43 +- server/src/kimi_k3/kimi_k3_internal.h | 6 +- .../test/test_moe_stream_owner_partition.cpp | 138 +++++++ variables.md | 4 +- 12 files changed, 835 insertions(+), 81 deletions(-) create mode 100644 server/test/test_moe_stream_owner_partition.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 9a8f94cf4..745b9717e 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1184,6 +1184,7 @@ if(DFLASH27B_TESTS) test/test_kvflash_pool_sizing.cpp test/test_kvflash_qk.cpp test/test_moe_expert_owner_placement.cpp + test/test_moe_stream_owner_partition.cpp test/test_qwen35moe_routing_stats.cpp test/test_qwen35moe_expert_placement.cpp test/test_qwen35moe_swap_manager.cpp diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index e0c1e1ba8..e9672c20a 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -28,7 +28,10 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_MMID_GROUPED_DEVICE` | -1 | Optional zero-based device restriction; unset/-1 applies to every eligible device. | | `DFLASH_DS4_MOE_TP` / `DFLASH_DS4_MOE_TP_INPROC` | unset | BURN-IN: enable DeepSeek4 route-owner expert parallelism in one process. | | `DFLASH_DS4_MOE_TP_GPU` | auto | HIP device for the cold DeepSeek4 expert owner. | -| `DFLASH_MOE_TP_GPU` | primary device | Model-neutral routed-expert owner GPU; used by Kimi streaming and preferred over legacy DS4/IPC GPU variables. | +| `DFLASH_MOE_TP_GPU` | primary device | BURN-IN: optional secondary routed-expert GPU; Kimi executes exact route partitions concurrently when it differs from the primary. | +| `DFLASH_MOE_PLACEMENT` | unset | BURN-IN: offline placement JSON; listed experts belong to the primary owner. | +| `DFLASH_MOE_PRIMARY_SHARE_PER_MILLE` | 500 | BURN-IN: deterministic primary route share when no explicit placement is supplied. | +| `DFLASH_MOE_DUAL_STREAM_TRACE` | unset | DEBUG: per-layer dual-owner route counts and branch/wall timings. | | `DFLASH_MOE_NVME_COLD_TIER` | auto | BURN-IN: DeepSeek4 cold-capacity policy (`auto`, `on`, `off`) for dual-device and Strix-only execution. | | `DFLASH_MOE_NVME_*` | tuned defaults | BURN-IN: bounded MoE SSD scheduler/backend controls; see `MOE_NVME_STREAMING.md`. | | `GGML_CUDA_BATCH_PEER_COPIES` | unset | BURN-IN: publish ordered HIP peer copies with one cross-device dependency per source/destination pair. | @@ -160,6 +163,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH_MODEL_CARDS_DIR` - model_card.cpp - `DFLASH_MOE_COLD_BACKEND` - deepseek4_loader.cpp - `DFLASH_MOE_COMPACT_MATERIALIZED` - moe_hybrid_ffn_eval.cpp +- `DFLASH_MOE_DUAL_STREAM_TRACE` - kimi_k3_graph.cpp - `DFLASH_MOE_DUPLICATE_HOT_ON_COLD` - moe_hybrid_storage.cpp - `DFLASH_MOE_EXPERT_COMPUTE_DAEMON_TOKEN_LOOP` - moe_expert_compute_ipc.cpp - `DFLASH_MOE_EXPERT_COMPUTE_IPC_BATCH_CAPACITY` - moe_expert_compute_ipc.cpp @@ -176,10 +180,12 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH_MOE_FIXED_SLOT_MAX` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_FULL_COLD_PARALLEL` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_FUSED_COMBINE` - moe_hybrid_ffn_eval.cpp +- `DFLASH_MOE_PLACEMENT` - kimi_k3_backend.cpp - `DFLASH_MOE_PREFILL_DEVICE_INPUT` - deepseek4_graph.cpp - `DFLASH_MOE_PREFILL_HOT_SUB_BATCH` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_PREFILL_MASKED_COLD` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_PREFILL_PERSISTENT_OWNER_ALLOC` - deepseek4_graph.cpp +- `DFLASH_MOE_PRIMARY_SHARE_PER_MILLE` - moe_hybrid_stream.cpp - `DFLASH_NO_MASK` - laguna_backend.cpp - `DFLASH_NO_MOE_ROUTER_FUSE` - qwen35moe_ffn.cpp - `DFLASH_NO_MOE_SWIGLU_FUSE` - qwen35moe_ffn.cpp diff --git a/server/docs/KIMI_K3_HETERO.md b/server/docs/KIMI_K3_HETERO.md index 978524981..1a178bee7 100644 --- a/server/docs/KIMI_K3_HETERO.md +++ b/server/docs/KIMI_K3_HETERO.md @@ -30,9 +30,13 @@ small architecture adapter around it: capped by the actual routed pool, so a tiny model cannot accidentally request a model-sized cache. - Routed-expert ownership is independent of contiguous layer splitting. - `DFLASH_MOE_TP_GPU` can place the SSD slots, adaptive cache, and selected - expert graphs on a second HIP device while KDA/MLA, recurrent/KV state, - shared experts, and sampling remain on the primary device. + With a second HIP device, two common stream engines partition the exact + native routes and run concurrently. The primary owns profile-selected hot + experts; the secondary owns the remaining capacity routes. Each engine has + independent SSD slots, adaptive cache, and persistent expert graphs. +- `MoeStreamDualOwnerExecutor` keeps one secondary worker alive for the model + lifetime. This removes thread create/join from every MoE layer and makes + routed wall time track the slower branch instead of their sum. - `bench_kimi_k3_hetero` runs Kimi's exact IQ1_S routed-expert geometry through the same model-neutral persistent evaluator used by production adapters: 896 experts, top-16, 92 MoE layers, @@ -49,9 +53,9 @@ The SSD text path is implemented, but two qualification boundaries remain: - The backend is correctness-first and token-sequential. Its per-layer graph boundaries are not yet fused/captured for full-model speed. -- The heterogeneous path currently crosses a host-visible, activation-sized - boundary at every routed layer. It is correct on two devices but does not - yet use PR-505's device-resident owner fork/join and overlap. +- The heterogeneous routed branches overlap, but their partial outputs still + cross a host-visible, activation-sized boundary at every routed layer. It + does not yet use PR-505's device-resident peer join. The vision encoder is out of scope for this text-only path. @@ -136,12 +140,13 @@ reserves. The implemented capacity-safe Strix-only starting point is: 2. NVMe: all routed expert stacks, with actual route misses read directly into pinned slots and evaluated on Strix. -This works unchanged on a Strix-only machine. A full Lucebox can additionally -assign the R9700 as the routed-expert owner with `DFLASH_MOE_TP_GPU`, moving -the adaptive expert cache and selected expert compute off Strix. That uses the -second device's memory, but the measured R9700 cold-storage path is slower, so -it is a capacity option rather than the default speed profile. The inverse -placement is also supported when a model's non-routed tensors fit on R9700. +This works unchanged on a Strix-only machine. When the non-routed execution +plan fits the R9700 budget, the full-Lucebox speed topology instead makes R9700 +the primary and sets `DFLASH_MOE_TP_GPU` to Strix. An offline placement +file assigns hot routes to R9700 while Strix executes the remaining routes and +serves the larger capacity cache. Both branches run concurrently; blindly +sending every streamed expert to R9700 remains slower because its cold SSD +upload path measured below Strix. After approximately 57.94 GiB of non-routed weights plus OS, workspace, and a moderate context reserve, roughly 70-85 GiB may remain for routed experts. @@ -158,8 +163,8 @@ can establish it. ## Next full-model milestone -The implementation no longer needs another generic cache or a second Kimi -router. Full-scale qualification requires: +The implementation no longer needs another generic cache, a second Kimi +router, or per-layer worker creation. Full-scale qualification requires: 1. Stage all 14 IQ1_S shards and run a short token-for-token comparison against the upstream Kimi implementation. @@ -168,8 +173,8 @@ router. Full-scale qualification requires: chosen context budget. 3. Compare end-to-end output and token rate against ordinary llama.cpp CPU/GPU offload. -4. Only after that baseline, replace the correctness-first host boundary with - a device-resident owner fork/join and overlap dense work with expert service. +4. Replace the correctness-first host partial join with a device-resident peer + join, then tune the placement until the two owner branches balance. The full 594 GB model cannot currently be staged on the qualification box, which currently has about 513 GB free. It needs at least about 650 GB of safe diff --git a/server/docs/MOE_NVME_STREAMING.md b/server/docs/MOE_NVME_STREAMING.md index 1f3111c3c..b2084da34 100644 --- a/server/docs/MOE_NVME_STREAMING.md +++ b/server/docs/MOE_NVME_STREAMING.md @@ -7,13 +7,21 @@ from SSD without changing their format or numerical representation. On a full Lucebox, the three owners have distinct jobs: -1. **R9700** owns dense layers and the statically hot routed experts. -2. **Strix Halo** owns an adaptive warm-expert cache, the bounded SSD staging - buffers, and execution of streamed cold experts. Its direct path to system - memory is faster than staging those experts into the discrete GPU on the - qualified machine. +1. **R9700** is the primary compute owner. It runs dense layers and the + profile-selected hot routed branch from its own SSD-backed expert cache. +2. **Strix Halo** is the capacity owner. It runs the remaining routed branch + from a separate adaptive cache and bounded SSD staging pipeline. Its direct + path to system memory is faster than staging every cold expert into the + discrete GPU on the qualified machine. 3. **NVMe** is the capacity tier for true cache misses that do not fit in the - Strix safe-memory budget. + two safe-memory budgets. + +The two routed branches launch concurrently. Each native route has exactly one +owner, and their activation-sized partials are added after both complete, so +the routed-layer target is `max(R9700 branch, Strix branch)` rather than their +sum. A persistent secondary worker avoids thread creation in the layer loop. +This route-owner split does not itself solve placement of a non-routed core +larger than R9700 VRAM; that remains an independent layer/tensor-plan decision. On a Strix Halo-only machine, the same device owns dense/static-hot weights, the warm cache, and streamed-expert execution. The planner reserves KV and @@ -43,6 +51,12 @@ and compute pipeline remain unchanged. Each file range can select a different shard, while single-file models use shard zero. See [KIMI_K3_HETERO.md](KIMI_K3_HETERO.md) for the 14-shard Kimi K3 qualification. +`MoeStreamDualOwnerExecutor` composes two such engines without adding a model +router. `MoeHybridPlacement` can select the primary owner's experts per layer; +otherwise a deterministic hash is only a bring-up fallback. Duplicate routes +to one expert retain one owner, and a partition test verifies that the two +weight masks reconstruct the original route batch exactly. + ## Scheduler `MoeNvmeScheduler` provides a bounded asynchronous data plane: @@ -142,22 +156,27 @@ Kimi's native router remains authoritative. The current text backend is correctness-first and sequential; captured per-layer graphs and the vision tower are separate optimizations. -On a two-GPU Lucebox, the primary device still owns KDA/MLA and model state; -the optional expert device owns the SSD slots, adaptive cache, and selected -expert graphs: +On a two-GPU Lucebox, put the compute-intensive primary path on the R9700 and +use Strix as the secondary capacity owner. Both devices receive independent +SSD slots/caches and execute their selected routed experts concurrently: ```bash -export DFLASH_MOE_TP_GPU= +export DFLASH_MOE_TP_GPU= ./build-hip-dual/dflash_server \ /path/to/Kimi-K3-UD-IQ1_S-00001-of-00014.gguf \ - --target-device hip: --max-ctx 8192 + --target-device hip: --max-ctx 8192 ``` This is functional expert ownership, not a contiguous layer split, so do not -use `--target-devices`. Leaving `DFLASH_MOE_TP_GPU` unset preserves the -Strix-only path. The current Kimi boundary uses activation-sized host staging; -device-resident peer fork/join is a later throughput optimization. +use `--target-devices`. `DFLASH_MOE_PLACEMENT` may point at an offline +`MoeHybridPlacement` JSON; its hot expert IDs become R9700-owned and +all other selected routes become Strix-owned. Without a plan, +`DFLASH_MOE_PRIMARY_SHARE_PER_MILLE` controls a deterministic bring-up split +(default 500). Leaving `DFLASH_MOE_TP_GPU` unset preserves the single-device +path, including Strix-only systems. The current Kimi join still uses +activation-sized host staging; a device-resident peer join remains the next +throughput optimization. ## Tuning and diagnostics @@ -177,7 +196,10 @@ not improve the qualified P310 drive and consumes extra pinned/system memory. | `DFLASH_MOE_NVME_DEVICE_CACHE_MB` | automatic/`0` | Override adaptive device expert-cache memory; `0` leaves only pipeline slots | | `DFLASH_MOE_NVME_GRAPH_CACHE` | `8` | Persistent expert-graph variants retained per stream engine; `0` is a diagnostic no-cache mode | | `DFLASH_MOE_NVME_REFERENCE_EVAL` | unset | Diagnostic only: `1` restores the allocation-heavy reference evaluator for numerical/performance A/B | -| `DFLASH_MOE_TP_GPU` | primary GPU | Optional second GPU that owns streamed expert cache and compute | +| `DFLASH_MOE_TP_GPU` | primary GPU | Optional secondary GPU; enables concurrent route ownership when different from the primary | +| `DFLASH_MOE_PLACEMENT` | unset | Offline placement JSON; listed experts belong to the primary GPU | +| `DFLASH_MOE_PRIMARY_SHARE_PER_MILLE` | `500` | Bring-up hash split used only when no placement is supplied | +| `DFLASH_MOE_DUAL_STREAM_TRACE` | unset | Debug per-layer owner counts and branch/wall timing | Shutdown telemetry reports logical and physical bytes, measured read service rate, cache hits, demand wait, de-duplication, dropped speculation, errors, and @@ -240,6 +262,15 @@ launches with one graph build, 167 graph-cache hits, and zero I/O errors. A 1 MiB device cache forced 163 evictions, demonstrating that the result came through the split-GGUF SSD path rather than accidental full residency. +The dual-owner extension was then qualified on the same fixture with the +R9700 primary and Strix secondary. Stable route ownership exercised both +engines and both issued real `io_uring` traffic; some top-2 layers naturally +landed on only one owner and bypassed the rendezvous. The eight greedy token +IDs matched the single-R9700 oracle exactly. After warm-up, dual-branch +routed-layer wall time matched the slower branch to within a few microseconds. +The tiny experts do not provide a meaningful end-to-end speed claim; full-size +expert geometry and a quiescent machine are required for that comparison. + ## Research lineage and next optimization The bounded priority/cache design follows the lessons of MoE-Infinity and @@ -249,9 +280,8 @@ that leaves the native router untouched. MoE-SpAc motivates compile-time expert layout and I/O coalescing. Tutti's slack-aware `io_uring` scheduling is relevant when persistent KV traffic shares the device. -The persistent graph and model-neutral evaluator are now shared by Kimi -qualification and production DS4 streaming. The next high-value work is a -repacked expert-major artifact, followed by overlap of hot-owner work with -cold-owner streaming. Any learned predictor comes after that deterministic -path is qualified, and may only prefetch routes selected later by the native -router. +The persistent graph, model-neutral evaluator, and exact concurrent owner +split are now shared infrastructure. The next high-value work is a repacked +expert-major artifact and a device-resident fork/join that removes Kimi's host +activation boundary. Any learned predictor comes after that deterministic path +is qualified, and may only prefetch routes selected later by the native router. diff --git a/server/src/common/moe_hybrid_stream.cpp b/server/src/common/moe_hybrid_stream.cpp index 1f78ba594..967554afa 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -6,13 +6,17 @@ #include "ggml-cuda.h" #include +#include #include +#include #include #include #include +#include #include #include #include +#include #include #include @@ -146,6 +150,14 @@ MoeStreamConfig MoeStreamConfig::from_env() { return config; } +MoeStreamDualOwnerPolicy MoeStreamDualOwnerPolicy::from_env() { + MoeStreamDualOwnerPolicy policy; + policy.primary_share_per_mille = env_bounded_int( + "DFLASH_MOE_PRIMARY_SHARE_PER_MILLE", + policy.primary_share_per_mille, 0, 1000); + return policy; +} + bool make_moe_stream_expert_spec( const MoeHybridConfig & cfg, const MoeLayerDesc & desc, @@ -1734,6 +1746,11 @@ bool eval_moe_streamed_experts( auto & runtime = *engine.runtime_; std::lock_guard compute_guard(runtime.compute_mutex); + ScopedGpuDevice device_scope(runtime.device); + if (!device_scope.ready()) { + if (err) *err = "failed to select streamed expert compute GPU"; + return false; + } if (!prepare_device_expert_layout( runtime, batch.layer, spec, err)) { return false; @@ -1968,6 +1985,367 @@ bool eval_moe_streamed_experts( return true; } +namespace { + +uint32_t stream_owner_hash(int layer, int expert) { + uint32_t value = (uint32_t) expert + 0x9e3779b9U; + value ^= (uint32_t) layer * 0x85ebca6bU; + value ^= value >> 16; + value *= 0x7feb352dU; + value ^= value >> 15; + value *= 0x846ca68bU; + value ^= value >> 16; + return value; +} + +} // namespace + +bool partition_moe_stream_routes( + const MoeStreamRouteBatch & batch, + const MoeStreamDualOwnerPolicy & policy, + std::vector & primary_weights, + std::vector & secondary_weights, + MoeStreamDualOwnerStats * stats, + std::string * err) { + if (batch.layer < 0 || batch.n_expert <= 0 || batch.top_k <= 0 || + batch.top_k > batch.n_expert || batch.n_tokens <= 0 || + !batch.selected_ids || !batch.selected_weights || + policy.primary_share_per_mille < 0 || + policy.primary_share_per_mille > 1000) { + if (err) *err = "invalid dual-owner route batch or policy"; + return false; + } + if (policy.primary_placement && + (policy.primary_placement->n_layer <= batch.layer || + policy.primary_placement->n_expert != batch.n_expert)) { + if (err) *err = "dual-owner placement does not match routed batch"; + return false; + } + + size_t route_slots = 0; + if (!checked_mul_size((size_t) batch.top_k, + (size_t) batch.n_tokens, route_slots)) { + if (err) *err = "dual-owner route slot count overflow"; + return false; + } + primary_weights.assign(route_slots, 0.0f); + secondary_weights.assign(route_slots, 0.0f); + + // -1 means unseen, 0 secondary, 1 primary. Duplicate appearances of one + // expert in a batch must retain one owner so its cache remains coherent. + std::vector owner((size_t) batch.n_expert, -1); + std::vector unique_experts; + unique_experts.reserve(std::min(route_slots, (size_t) batch.n_expert)); + int primary_experts = 0; + int secondary_experts = 0; + for (size_t route = 0; route < route_slots; ++route) { + const int32_t expert = batch.selected_ids[route]; + const float weight = batch.selected_weights[route]; + if (expert < 0 || weight == 0.0f) continue; + if (expert >= batch.n_expert || !std::isfinite(weight)) { + if (err) *err = "native router produced an invalid dual-owner route"; + return false; + } + if (owner[(size_t) expert] >= 0) continue; + bool primary_owner = false; + if (policy.primary_placement) { + primary_owner = policy.primary_placement->is_hot( + batch.layer, expert); + } else { + primary_owner = + stream_owner_hash(batch.layer, expert) % 1000U < + (uint32_t) policy.primary_share_per_mille; + } + owner[(size_t) expert] = primary_owner ? 1 : 0; + unique_experts.push_back(expert); + if (primary_owner) ++primary_experts; + else ++secondary_experts; + } + + int primary_routes = 0; + int secondary_routes = 0; + for (size_t route = 0; route < route_slots; ++route) { + const int32_t expert = batch.selected_ids[route]; + const float weight = batch.selected_weights[route]; + if (expert < 0 || weight == 0.0f) continue; + if (owner[(size_t) expert] == 1) { + primary_weights[route] = weight; + ++primary_routes; + } else { + secondary_weights[route] = weight; + ++secondary_routes; + } + } + if (stats) { + stats->primary_routes = primary_routes; + stats->secondary_routes = secondary_routes; + stats->primary_experts = primary_experts; + stats->secondary_experts = secondary_experts; + } + return true; +} + +struct MoeStreamDualOwnerExecutor::Runtime { + MoeHybridStreamEngine * primary = nullptr; + MoeHybridStreamEngine * secondary = nullptr; + std::mutex call_mutex; + std::mutex work_mutex; + std::condition_variable work_cv; + std::condition_variable done_cv; + std::thread worker; + bool stop = false; + bool pending = false; + bool done = false; + + const MoeStreamExpertSpec * job_spec = nullptr; + MoeStreamRouteBatch job_batch; + std::vector * job_out = nullptr; + std::string * job_error = nullptr; + bool job_ok = false; + uint64_t job_us = 0; + + void worker_loop() { + using Clock = std::chrono::steady_clock; + for (;;) { + const MoeStreamExpertSpec * spec = nullptr; + MoeStreamRouteBatch batch; + std::vector * out = nullptr; + std::string * error = nullptr; + { + std::unique_lock lock(work_mutex); + work_cv.wait(lock, [&]() { return stop || pending; }); + if (stop) return; + spec = job_spec; + batch = job_batch; + out = job_out; + error = job_error; + pending = false; + } + + const auto start = Clock::now(); + bool ok = false; + try { + ok = eval_moe_streamed_experts( + *secondary, *spec, batch, *out, error); + } catch (const std::exception & ex) { + *error = ex.what(); + } catch (...) { + *error = "unknown secondary-owner exception"; + } + const uint64_t elapsed_us = (uint64_t) + std::chrono::duration_cast( + Clock::now() - start).count(); + { + std::lock_guard lock(work_mutex); + job_ok = ok; + job_us = elapsed_us; + done = true; + } + done_cv.notify_one(); + } + } +}; + +MoeStreamDualOwnerExecutor::MoeStreamDualOwnerExecutor() = default; + +MoeStreamDualOwnerExecutor::~MoeStreamDualOwnerExecutor() { + destroy(); +} + +bool MoeStreamDualOwnerExecutor::init( + MoeHybridStreamEngine & primary, + MoeHybridStreamEngine & secondary, + std::string * err) { + destroy(); + if (!primary.is_bound() || !secondary.is_bound() || + !primary.compute_backend() || !secondary.compute_backend() || + primary.compute_backend() == secondary.compute_backend()) { + if (err) *err = "dual-owner streaming requires two bound GPU backends"; + return false; + } + + auto runtime = std::make_unique(); + runtime->primary = &primary; + runtime->secondary = &secondary; + try { + Runtime * worker_runtime = runtime.get(); + runtime->worker = std::thread( + [worker_runtime]() { worker_runtime->worker_loop(); }); + } catch (const std::exception & ex) { + if (err) { + *err = std::string("failed to start secondary owner worker: ") + + ex.what(); + } + return false; + } + runtime_ = std::move(runtime); + return true; +} + +bool MoeStreamDualOwnerExecutor::is_ready() const { + return runtime_ != nullptr; +} + +void MoeStreamDualOwnerExecutor::destroy() { + if (!runtime_) return; + auto runtime = std::move(runtime_); + std::lock_guard call_guard(runtime->call_mutex); + { + std::lock_guard lock(runtime->work_mutex); + runtime->stop = true; + } + runtime->work_cv.notify_one(); + if (runtime->worker.joinable()) runtime->worker.join(); +} + +bool MoeStreamDualOwnerExecutor::eval( + const MoeStreamExpertSpec & spec, + const MoeStreamRouteBatch & batch, + const MoeStreamDualOwnerPolicy & policy, + std::vector & out, + MoeStreamDualOwnerStats * stats, + std::string * err) { + if (!runtime_) { + if (err) *err = "dual-owner executor is not initialized"; + return false; + } + Runtime & runtime = *runtime_; + std::lock_guard call_guard(runtime.call_mutex); + if (!runtime.primary->is_bound() || !runtime.secondary->is_bound()) { + if (err) *err = "dual-owner stream engine was destroyed"; + return false; + } + + MoeStreamDualOwnerStats local_stats; + std::vector primary_weights; + std::vector secondary_weights; + if (!partition_moe_stream_routes( + batch, policy, primary_weights, secondary_weights, + &local_stats, err)) { + return false; + } + + // A stable placement may legitimately route this token to only one GPU. + // Bypass the worker rendezvous in that case; forcing synthetic work onto + // both owners would hurt cache locality and small-top-k decode latency. + if (local_stats.primary_routes == 0 || + local_stats.secondary_routes == 0) { + const bool use_primary = local_stats.primary_routes != 0; + MoeHybridStreamEngine & owner = use_primary + ? *runtime.primary : *runtime.secondary; + MoeStreamRouteBatch owner_batch = batch; + owner_batch.selected_weights = use_primary + ? primary_weights.data() : secondary_weights.data(); + using Clock = std::chrono::steady_clock; + const auto start = Clock::now(); + std::string owner_error; + bool ok = false; + try { + ok = eval_moe_streamed_experts( + owner, spec, owner_batch, out, &owner_error); + } catch (const std::exception & ex) { + owner_error = ex.what(); + } catch (...) { + owner_error = "unknown single-owner exception"; + } + const uint64_t elapsed_us = (uint64_t) + std::chrono::duration_cast( + Clock::now() - start).count(); + local_stats.wall_us = elapsed_us; + if (use_primary) local_stats.primary_us = elapsed_us; + else local_stats.secondary_us = elapsed_us; + if (stats) *stats = local_stats; + if (!ok && err) { + *err = std::string(use_primary ? "primary" : "secondary") + + " owner failed: " + owner_error; + } + return ok; + } + + MoeStreamRouteBatch primary_batch = batch; + MoeStreamRouteBatch secondary_batch = batch; + primary_batch.selected_weights = primary_weights.data(); + secondary_batch.selected_weights = secondary_weights.data(); + std::vector primary_out; + std::vector secondary_out; + std::string primary_error; + std::string secondary_error; + using Clock = std::chrono::steady_clock; + const auto wall_start = Clock::now(); + + { + std::lock_guard lock(runtime.work_mutex); + runtime.job_spec = &spec; + runtime.job_batch = secondary_batch; + runtime.job_out = &secondary_out; + runtime.job_error = &secondary_error; + runtime.job_ok = false; + runtime.job_us = 0; + runtime.done = false; + runtime.pending = true; + } + runtime.work_cv.notify_one(); + + bool primary_ok = false; + const auto primary_start = Clock::now(); + try { + primary_ok = eval_moe_streamed_experts( + *runtime.primary, spec, primary_batch, + primary_out, &primary_error); + } catch (const std::exception & ex) { + primary_error = ex.what(); + } catch (...) { + primary_error = "unknown primary-owner exception"; + } + local_stats.primary_us = (uint64_t) + std::chrono::duration_cast( + Clock::now() - primary_start).count(); + + bool secondary_ok = false; + { + std::unique_lock lock(runtime.work_mutex); + runtime.done_cv.wait(lock, [&]() { return runtime.done; }); + secondary_ok = runtime.job_ok; + local_stats.secondary_us = runtime.job_us; + } + local_stats.wall_us = (uint64_t) + std::chrono::duration_cast( + Clock::now() - wall_start).count(); + + if (!primary_ok || !secondary_ok) { + if (err) { + *err = !primary_ok + ? std::string("primary owner failed: ") + primary_error + : std::string("secondary owner failed: ") + secondary_error; + } + return false; + } + if (primary_out.size() != secondary_out.size()) { + if (err) *err = "dual-owner partial sizes do not match"; + return false; + } + out.resize(primary_out.size()); + for (size_t i = 0; i < out.size(); ++i) { + out[i] = primary_out[i] + secondary_out[i]; + } + if (stats) *stats = local_stats; + return true; +} + +bool eval_moe_streamed_experts_dual_owner( + MoeHybridStreamEngine & primary, + MoeHybridStreamEngine & secondary, + const MoeStreamExpertSpec & spec, + const MoeStreamRouteBatch & batch, + const MoeStreamDualOwnerPolicy & policy, + std::vector & out, + MoeStreamDualOwnerStats * stats, + std::string * err) { + MoeStreamDualOwnerExecutor executor; + if (!executor.init(primary, secondary, err)) return false; + return executor.eval(spec, batch, policy, out, stats, err); +} + bool eval_moe_cold_experts_streaming( MoeHybridStreamEngine & engine, ggml_backend_t gpu_backend, diff --git a/server/src/common/moe_hybrid_stream.h b/server/src/common/moe_hybrid_stream.h index 350606464..1c7ed1c3b 100644 --- a/server/src/common/moe_hybrid_stream.h +++ b/server/src/common/moe_hybrid_stream.h @@ -102,6 +102,27 @@ struct MoeStreamComputeStats { uint64_t graph_launches = 0; }; +// Route ownership for two concurrent SSD-backed GPU owners. An explicit +// placement takes precedence and identifies the primary GPU's hot experts. +// Without one, a stable layer/expert hash supplies a deterministic capacity +// split that preserves cache locality across tokens. +struct MoeStreamDualOwnerPolicy { + const MoeHybridPlacement * primary_placement = nullptr; + int primary_share_per_mille = 500; + + static MoeStreamDualOwnerPolicy from_env(); +}; + +struct MoeStreamDualOwnerStats { + uint64_t wall_us = 0; + uint64_t primary_us = 0; + uint64_t secondary_us = 0; + int primary_routes = 0; + int secondary_routes = 0; + int primary_experts = 0; + int secondary_experts = 0; +}; + class MoeHybridStreamEngine { public: MoeHybridStreamEngine(); @@ -198,6 +219,37 @@ class MoeHybridStreamEngine { std::unique_ptr runtime_; }; +// Persistent two-owner coordinator. The secondary compute worker is created +// once at model initialization, so decode does not pay a thread create/join at +// every routed layer. Calls are deliberately serialized; each call still +// launches the primary and secondary GPU pipelines concurrently. +class MoeStreamDualOwnerExecutor { +public: + MoeStreamDualOwnerExecutor(); + ~MoeStreamDualOwnerExecutor(); + + MoeStreamDualOwnerExecutor(const MoeStreamDualOwnerExecutor &) = delete; + MoeStreamDualOwnerExecutor & operator=( + const MoeStreamDualOwnerExecutor &) = delete; + + bool init(MoeHybridStreamEngine & primary, + MoeHybridStreamEngine & secondary, + std::string * err = nullptr); + bool is_ready() const; + void destroy(); + + bool eval(const MoeStreamExpertSpec & spec, + const MoeStreamRouteBatch & batch, + const MoeStreamDualOwnerPolicy & policy, + std::vector & out, + MoeStreamDualOwnerStats * stats = nullptr, + std::string * err = nullptr); + +private: + struct Runtime; + std::unique_ptr runtime_; +}; + // Evaluate exactly the experts selected by the native router. Placement only // decides which selected IDs arrive here; no prediction or cache policy can // change the returned mathematical function. @@ -208,6 +260,27 @@ bool eval_moe_streamed_experts( std::vector & out, std::string * err = nullptr); +// One-shot compatibility wrapper. Long-lived model adapters should initialize +// MoeStreamDualOwnerExecutor once to avoid per-layer thread startup overhead. +bool eval_moe_streamed_experts_dual_owner( + MoeHybridStreamEngine & primary, + MoeHybridStreamEngine & secondary, + const MoeStreamExpertSpec & spec, + const MoeStreamRouteBatch & batch, + const MoeStreamDualOwnerPolicy & policy, + std::vector & out, + MoeStreamDualOwnerStats * stats = nullptr, + std::string * err = nullptr); + +// Exposed for deterministic, GPU-free policy tests and offline plan tooling. +bool partition_moe_stream_routes( + const MoeStreamRouteBatch & batch, + const MoeStreamDualOwnerPolicy & policy, + std::vector & primary_weights, + std::vector & secondary_weights, + MoeStreamDualOwnerStats * stats = nullptr, + std::string * err = nullptr); + // Evaluate the cold contribution for one layer. All routed SSD requests are // admitted before compute starts, then double-buffered H2D runs concurrently // with the preceding expert graph. diff --git a/server/src/kimi_k3/kimi_k3_backend.cpp b/server/src/kimi_k3/kimi_k3_backend.cpp index 6b2ab1611..3e90eedff 100644 --- a/server/src/kimi_k3/kimi_k3_backend.cpp +++ b/server/src/kimi_k3/kimi_k3_backend.cpp @@ -72,15 +72,16 @@ bool KimiK3Backend::init_streaming() { return false; } } - ggml_backend_t stream_backend = - expert_backend_ ? expert_backend_ : backend_; auto fail_streaming = [&]() { + dual_stream_executor_.destroy(); stream_engine_.destroy(); + secondary_stream_engine_.destroy(); + stream_owner_policy_ = MoeStreamDualOwnerPolicy{}; + stream_placement_ = MoeHybridPlacement{}; release_expert_backend(); return false; }; - MoeStreamConfig stream_config = MoeStreamConfig::from_env(); size_t routed_pool_bytes = 0; for (const LayerExpertRegions & regions : weights_.streamed_layer_regions) { @@ -109,11 +110,16 @@ bool KimiK3Backend::init_streaming() { routed_pool_bytes += bytes_per_expert * static_cast(weights_.n_expert); } - if (!std::getenv("DFLASH_MOE_NVME_DEVICE_CACHE_MB")) { + auto stream_config_for = [&](int gpu, const char * owner_name) { + MoeStreamConfig stream_config = MoeStreamConfig::from_env(); + if (std::getenv("DFLASH_MOE_NVME_DEVICE_CACHE_MB")) { + stream_config.device_cache_bytes = + std::min(stream_config.device_cache_bytes, routed_pool_bytes); + return stream_config; + } size_t free_bytes = 0; size_t total_bytes = 0; - ggml_backend_cuda_get_device_memory( - expert_gpu_, &free_bytes, &total_bytes); + ggml_backend_cuda_get_device_memory(gpu, &free_bytes, &total_bytes); const size_t gib = 1024ULL * 1024ULL * 1024ULL; const size_t reserve = std::max(2 * gib, total_bytes / 20); stream_config.device_cache_bytes = @@ -121,24 +127,56 @@ bool KimiK3Backend::init_streaming() { ? std::min(free_bytes - reserve, routed_pool_bytes) : 0; std::fprintf(stderr, - "[kimi-k3] streamed expert cache: free=%.2f GiB reserve=%.2f GiB " - "pool=%.2f GiB cache=%.2f GiB\n", + "[kimi-k3] %s streamed cache: gpu=%d free=%.2f GiB " + "reserve=%.2f GiB pool=%.2f GiB cache=%.2f GiB\n", + owner_name, gpu, static_cast(free_bytes) / gib, static_cast(reserve) / gib, static_cast(routed_pool_bytes) / gib, static_cast(stream_config.device_cache_bytes) / gib); - } else { - stream_config.device_cache_bytes = - std::min(stream_config.device_cache_bytes, routed_pool_bytes); - } + return stream_config; + }; + + const MoeStreamConfig primary_config = stream_config_for( + cfg_.device.primary_gpu(), "primary"); if (!stream_engine_.init( - stream_backend, weights_.max_streamed_expert_bytes, - stream_config, &error)) { + backend_, weights_.max_streamed_expert_bytes, + primary_config, &error)) { std::fprintf(stderr, - "[kimi-k3] stream engine initialization failed: %s\n", + "[kimi-k3] primary stream engine initialization failed: %s\n", error.c_str()); return fail_streaming(); } + if (expert_backend_) { + const MoeStreamConfig secondary_config = stream_config_for( + expert_gpu_, "secondary"); + if (!secondary_stream_engine_.init( + expert_backend_, weights_.max_streamed_expert_bytes, + secondary_config, &error)) { + std::fprintf(stderr, + "[kimi-k3] secondary stream engine initialization failed: %s\n", + error.c_str()); + return fail_streaming(); + } + } + + stream_owner_policy_ = MoeStreamDualOwnerPolicy::from_env(); + stream_placement_ = MoeHybridPlacement{}; + const char * placement_path = std::getenv("DFLASH_MOE_PLACEMENT"); + if (expert_backend_ && placement_path && *placement_path) { + if (!MoeHybridPlacement::load_json( + placement_path, stream_placement_, &error) || + !stream_placement_.matches( + static_cast(weights_.streamed_layer_regions.size()), + weights_.n_expert, weights_.n_expert_used)) { + std::fprintf(stderr, + "[kimi-k3] invalid dual-owner placement %s: %s\n", + placement_path, + error.empty() ? "model shape mismatch" : error.c_str()); + return fail_streaming(); + } + stream_owner_policy_.primary_placement = &stream_placement_; + } std::vector descriptors; std::vector sources; @@ -198,8 +236,14 @@ bool KimiK3Backend::init_streaming() { sources.push_back({ nullptr, static_cast(shard_bytes), fd}); } - const bool bound = stream_engine_.bind_sources( + const bool primary_bound = stream_engine_.bind_sources( sources, weights_.streamed_layer_regions, &error); + bool secondary_bound = true; + std::string secondary_error; + if (primary_bound && expert_backend_) { + secondary_bound = secondary_stream_engine_.bind_sources( + sources, weights_.streamed_layer_regions, &secondary_error); + } for (int fd : descriptors) { #if defined(_WIN32) ::_close(fd); @@ -207,21 +251,45 @@ bool KimiK3Backend::init_streaming() { ::close(fd); #endif } - if (!bound) { + if (!primary_bound || !secondary_bound) { std::fprintf(stderr, "[kimi-k3] stream source binding failed: %s\n", + primary_bound ? secondary_error.c_str() : error.c_str()); + return fail_streaming(); + } + if (expert_backend_ && !dual_stream_executor_.init( + stream_engine_, secondary_stream_engine_, &error)) { + std::fprintf(stderr, + "[kimi-k3] dual-owner executor initialization failed: %s\n", error.c_str()); return fail_streaming(); } - std::fprintf(stderr, - "[kimi-k3] routed experts file-backed: shards=%zu layers=%zu " - "io=%s primary_gpu=%d expert_gpu=%d cache=%.2f GiB\n", - weights_.shard_paths.size(), - weights_.streamed_layer_regions.size(), - stream_engine_.io_backend_name(), - cfg_.device.primary_gpu(), expert_gpu_, - static_cast(stream_engine_.device_cache_bytes()) / - (1024.0 * 1024.0 * 1024.0)); + if (expert_backend_) { + std::fprintf(stderr, + "[kimi-k3] routed experts dual-owner: shards=%zu layers=%zu " + "primary=%d/%s/%.2fGiB secondary=%d/%s/%.2fGiB " + "primary_share=%d/1000 placement=%s\n", + weights_.shard_paths.size(), + weights_.streamed_layer_regions.size(), + cfg_.device.primary_gpu(), stream_engine_.io_backend_name(), + static_cast(stream_engine_.device_cache_bytes()) / + (1024.0 * 1024.0 * 1024.0), + expert_gpu_, secondary_stream_engine_.io_backend_name(), + static_cast(secondary_stream_engine_.device_cache_bytes()) / + (1024.0 * 1024.0 * 1024.0), + stream_owner_policy_.primary_share_per_mille, + stream_owner_policy_.primary_placement ? "profile" : "hash"); + } else { + std::fprintf(stderr, + "[kimi-k3] routed experts file-backed: shards=%zu layers=%zu " + "io=%s gpu=%d cache=%.2f GiB\n", + weights_.shard_paths.size(), + weights_.streamed_layer_regions.size(), + stream_engine_.io_backend_name(), + cfg_.device.primary_gpu(), + static_cast(stream_engine_.device_cache_bytes()) / + (1024.0 * 1024.0 * 1024.0)); + } return true; } @@ -252,11 +320,10 @@ bool KimiK3Backend::init() { if (weights_.routed_experts_streamed && !init_streaming()) return false; std::fprintf(stderr, "[kimi-k3] native backend ready on device %d (max_ctx=%d, " - "experts=%s:%d, correctness-first sequential prefill)\n", + "experts=%s, correctness-first sequential prefill)\n", cfg_.device.primary_gpu(), max_ctx, - weights_.routed_experts_streamed ? "nvme" : "resident", - weights_.routed_experts_streamed ? expert_gpu_ - : cfg_.device.primary_gpu()); + !weights_.routed_experts_streamed ? "resident" : + (expert_backend_ ? "nvme-dual-owner" : "nvme-single-owner")); std::fflush(stderr); return true; } @@ -272,7 +339,9 @@ void KimiK3Backend::print_ready_banner() const { bool KimiK3Backend::park(ParkTarget target) { if (!park_target_includes_target_model(target)) return false; if (!parked_) { + dual_stream_executor_.destroy(); stream_engine_.destroy(); + secondary_stream_engine_.destroy(); release_expert_backend(); free_kimi_k3_weights(weights_); parked_ = true; @@ -332,8 +401,12 @@ GenerateResult KimiK3Backend::generate_impl(const GenerateRequest & req, std::vector logits; const auto prefill_begin = std::chrono::steady_clock::now(); for (size_t i = 0; i < req.prompt.size(); ++i) { - if (!kimi_k3_step(backend_, weights_, cache_, req.prompt[i], - static_cast(i), logits, &stream_engine_)) { + if (!kimi_k3_step( + backend_, weights_, cache_, req.prompt[i], + static_cast(i), logits, &stream_engine_, + dual_stream_executor_.is_ready() + ? &dual_stream_executor_ : nullptr, + &stream_owner_policy_)) { result.fail(GenerateErrorCode::PrefillFailed, dflash27b_last_error()); out_io.emit(-1); @@ -374,8 +447,12 @@ GenerateResult KimiK3Backend::generate_impl(const GenerateRequest & req, out_io.emit(next); if (out_io.cancelled || next == weights_.eos_token_id) break; if (i + 1 < req.n_gen) { - if (!kimi_k3_step(backend_, weights_, cache_, next, - cache_.cur_pos, logits, &stream_engine_)) { + if (!kimi_k3_step( + backend_, weights_, cache_, next, + cache_.cur_pos, logits, &stream_engine_, + dual_stream_executor_.is_ready() + ? &dual_stream_executor_ : nullptr, + &stream_owner_policy_)) { result.fail(GenerateErrorCode::DecodeFailed, dflash27b_last_error()); out_io.emit(-1); @@ -428,7 +505,9 @@ bool KimiK3Backend::handle_compress(const std::string & line, } void KimiK3Backend::shutdown() { + dual_stream_executor_.destroy(); stream_engine_.destroy(); + secondary_stream_engine_.destroy(); release_expert_backend(); free_kimi_k3_cache(cache_); free_kimi_k3_weights(weights_); diff --git a/server/src/kimi_k3/kimi_k3_backend.h b/server/src/kimi_k3/kimi_k3_backend.h index 6f79cd493..2026f7b2e 100644 --- a/server/src/kimi_k3/kimi_k3_backend.h +++ b/server/src/kimi_k3/kimi_k3_backend.h @@ -15,8 +15,9 @@ struct KimiK3BackendConfig { DevicePlacement device; int stream_fd = -1; // -1 resolves DFLASH_MOE_TP_GPU and otherwise keeps experts on the - // primary GPU. A different GPU owns streamed weights/cache/compute while - // dense KDA/MLA, recurrent state, and sampling remain primary-owned. + // primary GPU. A different GPU becomes the secondary capacity owner; + // routed work is partitioned between both GPUs while dense KDA/MLA, + // recurrent state, and sampling remain primary-owned. int expert_gpu = -1; // Production Kimi uses file-backed routed experts. The resident mode is // retained as a deterministic oracle for small architecture fixtures. @@ -66,6 +67,10 @@ class KimiK3Backend final : public ModelBackend { KimiK3Weights weights_; KimiK3Cache cache_; MoeHybridStreamEngine stream_engine_; + MoeHybridStreamEngine secondary_stream_engine_; + MoeStreamDualOwnerExecutor dual_stream_executor_; + MoeHybridPlacement stream_placement_; + MoeStreamDualOwnerPolicy stream_owner_policy_; bool parked_ = false; std::mt19937_64 rng_{std::random_device{}()}; }; diff --git a/server/src/kimi_k3/kimi_k3_graph.cpp b/server/src/kimi_k3/kimi_k3_graph.cpp index da6e0810b..cc78dfd27 100644 --- a/server/src/kimi_k3/kimi_k3_graph.cpp +++ b/server/src/kimi_k3/kimi_k3_graph.cpp @@ -9,6 +9,8 @@ #include #include #include +#include +#include #include #include @@ -446,7 +448,9 @@ bool streamed_kimi_k3_step( int32_t token, int position, std::vector & logits, - MoeHybridStreamEngine & stream_engine) { + MoeHybridStreamEngine & stream_engine, + MoeStreamDualOwnerExecutor * dual_stream_executor, + const MoeStreamDualOwnerPolicy * stream_owner_policy) { std::vector hidden(static_cast(w.n_embd)); { @@ -637,9 +641,16 @@ bool streamed_kimi_k3_step( route_batch.selected_weights = route_weights.data(); std::vector routed_output; std::string stream_error; - if (!eval_moe_streamed_experts( + MoeStreamDualOwnerStats owner_stats; + const bool dual_owner = dual_stream_executor != nullptr; + const bool route_ok = dual_owner + ? dual_stream_executor->eval( + spec, route_batch, *stream_owner_policy, + routed_output, &owner_stats, &stream_error) + : eval_moe_streamed_experts( stream_engine, spec, route_batch, - routed_output, &stream_error)) { + routed_output, &stream_error); + if (!route_ok) { set_last_error( "Kimi-K3 routed layer " + std::to_string(il) + @@ -647,6 +658,18 @@ bool streamed_kimi_k3_step( stream_error); return false; } + const char * trace = std::getenv("DFLASH_MOE_DUAL_STREAM_TRACE"); + if (dual_owner && trace && *trace && std::strcmp(trace, "0") != 0) { + std::fprintf(stderr, + "[kimi-k3] dual-owner layer=%d routes=%d/%d experts=%d/%d " + "primary=%.3fms secondary=%.3fms wall=%.3fms\n", + route_batch.layer, + owner_stats.primary_routes, owner_stats.secondary_routes, + owner_stats.primary_experts, owner_stats.secondary_experts, + owner_stats.primary_us / 1000.0, + owner_stats.secondary_us / 1000.0, + owner_stats.wall_us / 1000.0); + } ctx = new_kimi_step_context(); if (!ctx) { @@ -797,7 +820,9 @@ bool kimi_k3_step(ggml_backend_t backend, int32_t token, int position, std::vector & logits, - MoeHybridStreamEngine * stream_engine) { + MoeHybridStreamEngine * stream_engine, + MoeStreamDualOwnerExecutor * dual_stream_executor, + const MoeStreamDualOwnerPolicy * stream_owner_policy) { if (!backend || !w.ctx || !cache.ctx || position < 0 || position >= cache.max_ctx || position != cache.cur_pos || token < 0 || token >= w.n_vocab) { @@ -810,9 +835,17 @@ bool kimi_k3_step(ggml_backend_t backend, "Kimi-K3 step: file-backed experts require a bound stream engine"); return false; } + if (dual_stream_executor && + (!dual_stream_executor->is_ready() || !stream_owner_policy)) { + set_last_error( + "Kimi-K3 step: dual-owner streaming requires a ready " + "executor and an ownership policy"); + return false; + } return streamed_kimi_k3_step( backend, w, cache, token, position, - logits, *stream_engine); + logits, *stream_engine, dual_stream_executor, + stream_owner_policy); } ggml_init_params params{}; diff --git a/server/src/kimi_k3/kimi_k3_internal.h b/server/src/kimi_k3/kimi_k3_internal.h index 5b82a983e..1f1f9c48f 100644 --- a/server/src/kimi_k3/kimi_k3_internal.h +++ b/server/src/kimi_k3/kimi_k3_internal.h @@ -24,6 +24,8 @@ namespace dflash::common { class MoeHybridStreamEngine; +class MoeStreamDualOwnerExecutor; +struct MoeStreamDualOwnerPolicy; struct KimiK3Layer { bool recurrent = false; @@ -172,6 +174,8 @@ bool kimi_k3_step(ggml_backend_t backend, int32_t token, int position, std::vector & logits, - MoeHybridStreamEngine * stream_engine = nullptr); + MoeHybridStreamEngine * stream_engine = nullptr, + MoeStreamDualOwnerExecutor * dual_stream_executor = nullptr, + const MoeStreamDualOwnerPolicy * stream_owner_policy = nullptr); } // namespace dflash::common diff --git a/server/test/test_moe_stream_owner_partition.cpp b/server/test/test_moe_stream_owner_partition.cpp new file mode 100644 index 000000000..52e1ca1fc --- /dev/null +++ b/server/test/test_moe_stream_owner_partition.cpp @@ -0,0 +1,138 @@ +#include "CppUnitTestFramework.hpp" +#include "../src/common/moe_hybrid_stream.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +namespace { + +struct MoeStreamOwnerPartitionFixture {}; + +MoeStreamRouteBatch make_batch(const std::vector & ids, + const std::vector & weights) { + MoeStreamRouteBatch batch; + batch.layer = 0; + batch.n_expert = 8; + batch.top_k = 4; + batch.n_tokens = 2; + batch.selected_ids = ids.data(); + batch.selected_weights = weights.data(); + return batch; +} + +bool is_exact_partition(const std::vector & weights, + const std::vector & primary, + const std::vector & secondary) { + if (primary.size() != weights.size() || + secondary.size() != weights.size()) return false; + for (size_t i = 0; i < weights.size(); ++i) { + if (primary[i] + secondary[i] != weights[i] || + (primary[i] != 0.0f && secondary[i] != 0.0f)) return false; + } + return true; +} + +} // namespace + +TEST_CASE(MoeStreamOwnerPartitionFixture, + hash_policy_is_exact_deterministic_and_uses_both_owners) { + const std::vector ids = {0, 1, 2, 3, 1, 4, 5, 6}; + const std::vector weights = { + 0.40f, 0.30f, 0.20f, 0.10f, + 0.35f, 0.30f, 0.20f, 0.15f}; + const MoeStreamRouteBatch batch = make_batch(ids, weights); + MoeStreamDualOwnerPolicy policy; + policy.primary_share_per_mille = 500; + + std::vector primary; + std::vector secondary; + MoeStreamDualOwnerStats stats; + std::string error; + REQUIRE(partition_moe_stream_routes( + batch, policy, primary, secondary, &stats, &error)); + REQUIRE(is_exact_partition(weights, primary, secondary)); + REQUIRE(stats.primary_experts > 0); + REQUIRE(stats.secondary_experts > 0); + + // Expert 1 appears in both tokens and must retain one cache owner. + REQUIRE((primary[1] != 0.0f) == (primary[4] != 0.0f)); + + std::vector primary_again; + std::vector secondary_again; + REQUIRE(partition_moe_stream_routes( + batch, policy, primary_again, secondary_again, nullptr, &error)); + REQUIRE(primary_again == primary); + REQUIRE(secondary_again == secondary); +} + +TEST_CASE(MoeStreamOwnerPartitionFixture, + explicit_placement_is_authoritative) { + const std::vector ids = {0, 1, 2, 3, 1, 4, 5, 6}; + const std::vector weights = { + 0.40f, 0.30f, 0.20f, 0.10f, + 0.35f, 0.30f, 0.20f, 0.15f}; + const MoeStreamRouteBatch batch = make_batch(ids, weights); + + MoeHybridPlacement placement; + placement.n_layer = 1; + placement.n_expert = 8; + placement.n_expert_used = 4; + placement.total_hot = 2; + placement.hot_counts = {2}; + placement.hot_expert_ids = {{1, 5}}; + + MoeStreamDualOwnerPolicy policy; + policy.primary_placement = &placement; + policy.primary_share_per_mille = 500; + + std::vector primary; + std::vector secondary; + MoeStreamDualOwnerStats stats; + std::string error; + REQUIRE(partition_moe_stream_routes( + batch, policy, primary, secondary, &stats, &error)); + REQUIRE(is_exact_partition(weights, primary, secondary)); + for (size_t i = 0; i < ids.size(); ++i) { + const bool expected_primary = ids[i] == 1 || ids[i] == 5; + REQUIRE((primary[i] != 0.0f) == expected_primary); + } + REQUIRE(stats.primary_experts == 2); + REQUIRE(stats.secondary_experts == 5); +} + +TEST_CASE(MoeStreamOwnerPartitionFixture, + endpoint_shares_and_shape_mismatch_fail_closed) { + const std::vector ids = {0, 1, 2, 3, 1, 4, 5, 6}; + const std::vector weights(ids.size(), 0.25f); + const MoeStreamRouteBatch batch = make_batch(ids, weights); + std::vector primary; + std::vector secondary; + MoeStreamDualOwnerStats stats; + std::string error; + + MoeStreamDualOwnerPolicy policy; + policy.primary_share_per_mille = 0; + REQUIRE(partition_moe_stream_routes( + batch, policy, primary, secondary, &stats, &error)); + REQUIRE(stats.primary_routes == 0); + REQUIRE(stats.secondary_routes == 8); + REQUIRE(is_exact_partition(weights, primary, secondary)); + + policy.primary_share_per_mille = 1000; + REQUIRE(partition_moe_stream_routes( + batch, policy, primary, secondary, &stats, &error)); + REQUIRE(stats.primary_routes == 8); + REQUIRE(stats.secondary_routes == 0); + REQUIRE(is_exact_partition(weights, primary, secondary)); + + MoeHybridPlacement mismatch; + mismatch.n_layer = 1; + mismatch.n_expert = 7; + policy.primary_placement = &mismatch; + REQUIRE(!partition_moe_stream_routes( + batch, policy, primary, secondary, nullptr, &error)); +} diff --git a/variables.md b/variables.md index 822097ff7..3b4583bd5 100644 --- a/variables.md +++ b/variables.md @@ -217,7 +217,9 @@ Untagged variables are operational tuning knobs. | Variable | Purpose | |---|---| -| `DFLASH_MOE_TP_GPU` | Model-neutral HIP device that owns routed-expert cache and compute; Kimi uses the primary device when unset. | +| `DFLASH_MOE_TP_GPU` | Optional secondary HIP route owner; Kimi runs primary and secondary SSD-backed route partitions concurrently when it differs from the primary device. | +| `DFLASH_MOE_PLACEMENT` / `DFLASH_MOE_PRIMARY_SHARE_PER_MILLE` | Profile-selected primary experts, or deterministic bring-up share when no placement exists. | +| `DFLASH_MOE_DUAL_STREAM_TRACE` | Debug per-layer route ownership and concurrent branch timing. | | `DFLASH_MOE_HYBRID_PREFILL_EAGER` / `DFLASH_MOE_PREFILL_TRACE` | Model-neutral heterogeneous prefill policy and tracing. The legacy `DFLASH_DS4_*` spellings remain aliases. | | `DFLASH_MOE_TP_GROUPED_MMVQ` / `DFLASH_MOE_TP_FUSED_GATE_UP` | Model-neutral grouped and fused routed-FFN kernel qualification switches. The legacy `DFLASH_DS4_*` spellings remain aliases. | | `DFLASH_MOE_TP_COARSE_OWNER` / `DFLASH_MOE_TP_COARSE_OWNER_SPLIT` | Model-neutral owner-op lowering switches. The legacy `DFLASH_DS4_*` spellings remain aliases. | From 49a898c1627483a36728c9edf7a1033bd1b5ca54 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:34:27 +0200 Subject: [PATCH 11/20] fix(moe): harden NVMe streaming for production --- .github/workflows/ci.yml | 15 ++-- server/docs/MOE_NVME_STREAMING.md | 8 +- server/src/common/moe_hybrid_stream.cpp | 4 +- server/src/common/moe_nvme_scheduler.cpp | 108 ++++++++++++++++++---- server/src/common/moe_nvme_scheduler.h | 8 ++ server/test/test_moe_nvme_scheduler.cpp | 110 +++++++++++++++++++++++ 6 files changed, 228 insertions(+), 25 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 350960916..00804a9ab 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -79,12 +79,14 @@ jobs: -DCMAKE_BUILD_TYPE=Release cmake --build build --target \ test_dflash test_generate test_flash_attn_sparse test_server_unit \ - test_deepseek4_unit -j$(nproc) + test_deepseek4_unit test_moe_nvme_scheduler -j$(nproc) - name: Run C++ server unit tests run: | cd server/build - ctest --output-on-failure -R "server_unit|deepseek4_unit" --no-tests=error + ctest --output-on-failure \ + -R "server_unit|deepseek4_unit|test_moe_nvme_scheduler" \ + --no-tests=error - name: Populate venv with cu128 torch + setuptools # First pass: install the workspace's default deps. dflash declares @@ -174,7 +176,7 @@ jobs: -DCMAKE_BUILD_TYPE=Release cmake --build build \ --target test_flash_attn_sparse test_deepseek4_mmid_grouped_cuda \ - test_moe_stream_compute \ + test_moe_stream_compute test_moe_nvme_scheduler \ -j"$(nproc)" - name: Run flash-attn sparse kernel test on the 3090 @@ -188,6 +190,9 @@ jobs: - name: Run compact-to-padded MXFP4 expert streaming test on the 3090 run: ./server/build/test_moe_stream_compute + - name: Run NVMe scheduler production tests + run: ./server/build/test_moe_nvme_scheduler + # Optional model-backed end-to-end smoke (real spec-decode on the 3090), # disabled by default because it builds dflash_server and lazy-loads the # ~16 GB Qwen3.6-27B target + draft (~1-2 min). The weights are already @@ -295,11 +300,11 @@ jobs: cmake --build "$RUNNER_TEMP/rocmfp-build" \ --target test_rocmfp4 test_rocmfpx test_rocmfp4_hip_tail test_rocmfpx_mmq \ test_deepseek4_mmid_grouped_cuda test_moe_stream_compute \ - test_recurrent_snapshot test_server_unit \ + test_moe_nvme_scheduler test_recurrent_snapshot test_server_unit \ --parallel 8 ctest --test-dir "$RUNNER_TEMP/rocmfp-build" \ --output-on-failure \ - -R 'rocmfp4_reference|rocmfpx_reference|rocmfp4_hip_tail|rocmfpx_mmq|deepseek4_mmid_grouped_cuda|test_moe_stream_compute|recurrent_snapshot|ChainRollbackPolicy' + -R 'rocmfp4_reference|rocmfpx_reference|rocmfp4_hip_tail|rocmfpx_mmq|deepseek4_mmid_grouped_cuda|test_moe_stream_compute|test_moe_nvme_scheduler|recurrent_snapshot|ChainRollbackPolicy' build-windows: name: Build Windows (MSVC + CUDA, library + server targets) diff --git a/server/docs/MOE_NVME_STREAMING.md b/server/docs/MOE_NVME_STREAMING.md index b2084da34..47c2d50b0 100644 --- a/server/docs/MOE_NVME_STREAMING.md +++ b/server/docs/MOE_NVME_STREAMING.md @@ -192,6 +192,7 @@ not improve the qualified P310 drive and consumes extra pinned/system memory. | `DFLASH_MOE_NVME_IO_THREADS` | `4` | Portable pread workers | | `DFLASH_MOE_NVME_DEMAND_RESERVE` | `2` | Slots unavailable to speculation | | `DFLASH_MOE_NVME_PREFETCH_BATCH` | `2` | Maximum speculative jobs per ring submission | +| `DFLASH_MOE_NVME_DEMAND_TIMEOUT_MS` | `30000` | Maximum wait for a demanded expert; `0` disables the guard | | `DFLASH_MOE_NVME_DEVICE_SLOTS` | `2` | Minimum rotating GPU expert buffers | | `DFLASH_MOE_NVME_DEVICE_CACHE_MB` | automatic/`0` | Override adaptive device expert-cache memory; `0` leaves only pipeline slots | | `DFLASH_MOE_NVME_GRAPH_CACHE` | `8` | Persistent expert-graph variants retained per stream engine; `0` is a diagnostic no-cache mode | @@ -202,8 +203,11 @@ not improve the qualified P310 drive and consumes extra pinned/system memory. | `DFLASH_MOE_DUAL_STREAM_TRACE` | unset | Debug per-layer owner counts and branch/wall timing | Shutdown telemetry reports logical and physical bytes, measured read service -rate, cache hits, demand wait, de-duplication, dropped speculation, errors, and -persistent-graph builds/hits/evictions. `test_moe_stream_compute` generates +rate, cache hits, demand wait/timeouts, de-duplication, dropped speculation, +errors, and persistent-graph builds/hits/evictions. The scheduler rejects +truncated shards at bind time and accepts a valid short direct-I/O completion +only when it covers the complete logical payload at an unaligned file tail. +`test_moe_stream_compute` generates tiny experts and checks both tensor-major and expert-major GPU results against a CPU oracle. It defaults to GPU 0, so it runs directly on Strix-only systems; `DFLASH_TEST_GPU` selects another device on multi-GPU hosts. The standalone diff --git a/server/src/common/moe_hybrid_stream.cpp b/server/src/common/moe_hybrid_stream.cpp index 967554afa..72232edbf 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -969,7 +969,8 @@ void MoeHybridStreamEngine::destroy() { "[moe-nvme] io=%s requests=%llu reads=%llu " "payload=%.3f GiB physical=%.3f GiB active-io-rate=%.3f GiB/s " "cache-hit=%.1f%% mean-demand-wait=%.3f ms " - "dedupe=%llu upgrades=%llu dropped-prefetch=%llu errors=%llu " + "dedupe=%llu upgrades=%llu dropped-prefetch=%llu " + "timeouts=%llu errors=%llu " "device-cache=%.1f MiB slots=%zu hits=%llu misses=%llu evictions=%llu " "graphs=%llu graph-hits=%llu graph-evictions=%llu launches=%llu\n", runtime_->io->effective_backend_name(), @@ -979,6 +980,7 @@ void MoeHybridStreamEngine::destroy() { (unsigned long long) stats.inflight_deduplications, (unsigned long long) stats.demand_upgrades, (unsigned long long) stats.prefetch_drops, + (unsigned long long) stats.demand_timeouts, (unsigned long long) stats.errors, device_cache_byte_count / 1024.0 / 1024.0, device_cache_slot_count, diff --git a/server/src/common/moe_nvme_scheduler.cpp b/server/src/common/moe_nvme_scheduler.cpp index 540781d66..8b46cfd14 100644 --- a/server/src/common/moe_nvme_scheduler.cpp +++ b/server/src/common/moe_nvme_scheduler.cpp @@ -20,6 +20,7 @@ #if defined(_WIN32) #include +#include #else #include #include @@ -126,10 +127,17 @@ uint64_t physical_memory_bytes() { } #if !defined(_WIN32) -bool pread_full(int fd, uint8_t * dst, size_t bytes, size_t offset, std::string & err) { +bool pread_at_least(int fd, uint8_t * dst, size_t request_bytes, + size_t required_bytes, size_t offset, + size_t & bytes_read, std::string & err) { + bytes_read = 0; + if (required_bytes > request_bytes) { + err = "invalid minimum read length"; + return false; + } size_t done = 0; - while (done < bytes) { - const size_t remaining = bytes - done; + while (done < required_bytes) { + const size_t remaining = request_bytes - done; const size_t chunk = std::min(remaining, (size_t) std::numeric_limits::max()); const ssize_t got = ::pread(fd, dst + done, chunk, (off_t) (offset + done)); if (got < 0) { @@ -143,6 +151,7 @@ bool pread_full(int fd, uint8_t * dst, size_t bytes, size_t offset, std::string } done += (size_t) got; } + bytes_read = done; return true; } #endif @@ -380,6 +389,8 @@ MoeNvmeConfig MoeNvmeConfig::from_env(MoeNvmeConfig base) { "DFLASH_MOE_NVME_DEMAND_RESERVE", base.demand_reserve, 1, base.host_slots - 1); base.max_prefetch_batch = parse_bounded_int( "DFLASH_MOE_NVME_PREFETCH_BATCH", base.max_prefetch_batch, 1, base.host_slots); + base.demand_timeout_ms = parse_bounded_int( + "DFLASH_MOE_NVME_DEMAND_TIMEOUT_MS", base.demand_timeout_ms, 0, 600000); if (const char * value = std::getenv("DFLASH_MOE_NVME_BACKEND")) { const std::string mode = lowercase(value); @@ -498,6 +509,7 @@ bool make_moe_expert_io_layout( span.io_file_offset = aligned_file; span.io_buffer_offset = aligned_buffer; span.io_bytes = aligned_io_bytes; + span.io_required_bytes = raw_io_bytes; return true; }; @@ -668,6 +680,7 @@ struct MoeNvmeScheduler::Impl { uint64_t active_io_ns = 0; std::atomic read_ns{0}; std::atomic wait_ns{0}; + std::atomic demand_timeouts{0}; std::atomic errors{0}; void begin_io_activity() { @@ -791,6 +804,10 @@ struct MoeNvmeScheduler::Impl { Admission admit_locked(int layer, int expert, MoeNvmePriority priority, int & slot_out, std::string * err) { slot_out = -1; + if (stopping) { + if (err) *err = "SSD scheduler is stopping"; + return Admission::Invalid; + } if (!bound) { if (err) *err = "SSD scheduler has no bound model source"; return Admission::Invalid; @@ -905,9 +922,16 @@ struct MoeNvmeScheduler::Impl { #else const size_t read_offset = direct_active ? span.io_file_offset : span.file_offset; const size_t read_bytes = direct_active ? span.io_bytes : span.bytes; + const size_t required_bytes = direct_active + ? span.io_required_bytes : span.bytes; const size_t buffer_offset = direct_active ? span.io_buffer_offset : span.buffer_offset; - if (!pread_full(active_fd, base + buffer_offset, read_bytes, read_offset, err)) return false; - physical += read_bytes; + size_t actual_bytes = 0; + if (!pread_at_least(active_fd, base + buffer_offset, + read_bytes, required_bytes, read_offset, + actual_bytes, err)) { + return false; + } + physical += actual_bytes; #endif } } @@ -969,6 +993,7 @@ struct MoeNvmeScheduler::Impl { struct Op { size_t job = 0; uint32_t expected = 0; + uint32_t required = 0; }; struct Progress { int pending = 0; @@ -1018,11 +1043,14 @@ struct MoeNvmeScheduler::Impl { const MoeExpertIoSpan & span = slot.layout.spans[s]; const size_t read_offset = direct_active ? span.io_file_offset : span.file_offset; const size_t read_bytes = direct_active ? span.io_bytes : span.bytes; + const size_t required_bytes = direct_active + ? span.io_required_bytes : span.bytes; const size_t buffer_offset = direct_active ? span.io_buffer_offset : span.buffer_offset; - if (read_bytes > std::numeric_limits::max()) { + if (read_bytes > (size_t) std::numeric_limits::max() || + required_bytes > read_bytes) { progress[j].ok = false; progress[j].error = - "one expert tensor read exceeds io_uring's 32-bit length"; + "one expert tensor read exceeds io_uring's supported length"; continue; } io_uring_sqe * sqe = ring->get_sqe(); @@ -1032,14 +1060,14 @@ struct MoeNvmeScheduler::Impl { continue; } const uint64_t op_index = operations.size(); - operations.push_back({j, (uint32_t) read_bytes}); + operations.push_back({j, (uint32_t) read_bytes, + (uint32_t) required_bytes}); ring->prepare_read(sqe, active_fds, span.source_index, jobs[j].slot, base + buffer_offset, (uint32_t) read_bytes, read_offset, op_index, !direct_active); ++progress[j].pending; ++progress[j].ops; - progress[j].physical += read_bytes; } } @@ -1097,13 +1125,16 @@ struct MoeNvmeScheduler::Impl { } const Op & op = operations[(size_t) cqe.user_data]; Progress & item = progress[op.job]; - if (cqe.res != (int32_t) op.expected) { + if (cqe.res < 0) { item.ok = false; - if (cqe.res < 0) { - item.error = std::string("io_uring read failed: ") + - std::strerror(-cqe.res); - } else { - item.error = "io_uring returned a short model read"; + item.error = std::string("io_uring read failed: ") + + std::strerror(-cqe.res); + } else { + item.physical += (uint32_t) cqe.res; + if ((uint32_t) cqe.res < op.required || + (uint32_t) cqe.res > op.expected) { + item.ok = false; + item.error = "io_uring returned an incomplete model read"; } } if (item.pending > 0 && --item.pending == 0) finish_job(op.job); @@ -1191,6 +1222,7 @@ bool MoeNvmeScheduler::init(const MoeNvmeConfig & requested, p.config.host_slots - 1)); p.config.max_prefetch_batch = std::max(1, std::min(p.config.max_prefetch_batch, p.config.host_slots)); + p.config.demand_timeout_ms = std::max(0, p.config.demand_timeout_ms); if (!is_power_of_two(p.config.direct_alignment) || p.config.direct_alignment < 512 || max_expert_payload_bytes == 0 || !allocate || !free_fn) { if (err) *err = "invalid SSD scheduler initialization arguments"; @@ -1263,6 +1295,25 @@ bool MoeNvmeScheduler::bind_sources( } all_mapped = all_mapped && source.mmap_data != nullptr; all_have_fds = all_have_fds && source.fd >= 0; +#if defined(_WIN32) + if (source.fd >= 0) { + struct _stat64 file_stat{}; + if (::_fstat64(source.fd, &file_stat) != 0 || file_stat.st_size < 0 || + (uint64_t) file_stat.st_size < (uint64_t) source.mmap_size) { + if (err) *err = "model shard fd is unreadable or shorter than its declared size"; + return false; + } + } +#else + if (source.fd >= 0) { + struct stat file_stat{}; + if (::fstat(source.fd, &file_stat) != 0 || file_stat.st_size < 0 || + (uint64_t) file_stat.st_size < (uint64_t) source.mmap_size) { + if (err) *err = "model shard fd is unreadable or shorter than its declared size"; + return false; + } + } +#endif if ((uint64_t) source.mmap_size > std::numeric_limits::max() - total_source_bytes) { if (err) *err = "SSD model shard sizes overflow"; @@ -1545,11 +1596,32 @@ bool MoeNvmeScheduler::acquire(int layer, int expert, MoeNvmeLease & out, } Impl & p = *impl_; const auto begin = Clock::now(); + const bool timeout_enabled = p.config.demand_timeout_ms > 0; + const auto deadline = timeout_enabled + ? begin + std::chrono::milliseconds(p.config.demand_timeout_ms) + : Clock::time_point::max(); p.requests.fetch_add(1, std::memory_order_relaxed); p.demand_requests.fetch_add(1, std::memory_order_relaxed); const MoeExpertKey key{(int32_t) layer, (int32_t) expert}; std::unique_lock lock(p.mutex); + auto wait_for_state = [&]() -> bool { + if (!timeout_enabled) { + p.state_cv.wait(lock); + return true; + } + if (p.state_cv.wait_until(lock, deadline) != std::cv_status::timeout) { + return true; + } + p.demand_timeouts.fetch_add(1, std::memory_order_relaxed); + p.wait_ns.fetch_add(elapsed_ns(begin, Clock::now()), + std::memory_order_relaxed); + if (err) { + *err = "timed out waiting for an SSD expert after " + + std::to_string(p.config.demand_timeout_ms) + " ms"; + } + return false; + }; bool admitted = false; for (;;) { if (p.stopping) { @@ -1563,7 +1635,7 @@ bool MoeNvmeScheduler::acquire(int layer, int expert, MoeNvmeLease & out, layer, expert, MoeNvmePriority::Demand, slot, err); if (result == Impl::Admission::Invalid) return false; if (result == Impl::Admission::NoSlot) { - p.state_cv.wait(lock); + if (!wait_for_state()) return false; continue; } admitted = true; @@ -1607,7 +1679,7 @@ bool MoeNvmeScheduler::acquire(int layer, int expert, MoeNvmeLease & out, return false; } if (slot.state != Impl::SlotState::Ready) { - p.state_cv.wait(lock); + if (!wait_for_state()) return false; continue; } @@ -1654,6 +1726,7 @@ MoeNvmeStats MoeNvmeScheduler::stats() const { out.active_io_ns = p.active_io_time(); out.read_ns = p.read_ns.load(std::memory_order_relaxed); out.wait_ns = p.wait_ns.load(std::memory_order_relaxed); + out.demand_timeouts = p.demand_timeouts.load(std::memory_order_relaxed); out.errors = p.errors.load(std::memory_order_relaxed); return out; } @@ -1675,6 +1748,7 @@ void MoeNvmeScheduler::reset_stats() { p.reset_io_time(); p.read_ns.store(0, std::memory_order_relaxed); p.wait_ns.store(0, std::memory_order_relaxed); + p.demand_timeouts.store(0, std::memory_order_relaxed); p.errors.store(0, std::memory_order_relaxed); } diff --git a/server/src/common/moe_nvme_scheduler.h b/server/src/common/moe_nvme_scheduler.h index f803c3c1d..d931a4180 100644 --- a/server/src/common/moe_nvme_scheduler.h +++ b/server/src/common/moe_nvme_scheduler.h @@ -53,6 +53,10 @@ struct MoeNvmeConfig { // latency of a demand arriving immediately after speculation was issued. int max_prefetch_batch = 2; + // Bound a demand wait so a failed drive cannot hang an inference worker + // forever. Zero disables the timeout for diagnostic use. + int demand_timeout_ms = 30000; + size_t direct_alignment = 4096; // Environment overrides use the DFLASH_MOE_NVME_* prefix. Invalid values @@ -87,6 +91,9 @@ struct MoeExpertIoSpan { size_t io_file_offset = 0; // aligned direct-I/O start size_t io_buffer_offset = 0; // aligned direct-I/O destination size_t io_bytes = 0; // aligned direct-I/O length + // Bytes that must be returned to cover the logical payload. This may be + // smaller than io_bytes for the final unaligned page of a model shard. + size_t io_required_bytes = 0; }; enum class MoeExpertComponentKind : uint8_t { @@ -163,6 +170,7 @@ struct MoeNvmeStats { uint64_t active_io_ns = 0; uint64_t read_ns = 0; uint64_t wait_ns = 0; + uint64_t demand_timeouts = 0; uint64_t errors = 0; }; diff --git a/server/test/test_moe_nvme_scheduler.cpp b/server/test/test_moe_nvme_scheduler.cpp index 75f72afd0..ccb90356f 100644 --- a/server/test/test_moe_nvme_scheduler.cpp +++ b/server/test/test_moe_nvme_scheduler.cpp @@ -282,6 +282,34 @@ TEST_CASE(MoeNvmeSchedulerFixture, speculation_cannot_consume_demand_reserve) { NVME_REQUIRE(stats.errors == 0); } +TEST_CASE(MoeNvmeSchedulerFixture, demand_timeout_prevents_busy_cache_deadlock) { + SyntheticModel model; + MoeNvmeConfig config; + config.backend = MoeNvmeBackend::Mmap; + config.direct_io = MoeNvmeDirectMode::Disabled; + config.host_slots = 2; + config.io_threads = 1; + config.demand_timeout_ms = 25; + + MoeNvmeScheduler scheduler; + std::string err; + NVME_REQUIRE(scheduler.init(config, 32 * 1024, aligned_allocate, + aligned_free, nullptr, &err)); + NVME_REQUIRE(scheduler.bind_source( + {model.file.data(), model.file.size(), -1}, model.regions, &err)); + + MoeNvmeLease first; + MoeNvmeLease second; + MoeNvmeLease blocked; + NVME_REQUIRE(scheduler.acquire(0, 0, first, &err)); + NVME_REQUIRE(scheduler.acquire(0, 1, second, &err)); + NVME_REQUIRE(!scheduler.acquire(0, 2, blocked, &err)); + NVME_REQUIRE(err.find("timed out") != std::string::npos); + NVME_REQUIRE(scheduler.stats().demand_timeouts == 1); + first.reset(); + second.reset(); +} + TEST_CASE(MoeNvmeSchedulerFixture, split_model_reads_tensor_spans_from_multiple_shards) { constexpr int experts = SyntheticModel::kExperts; std::vector shard_a(512 * 1024, 0xa5); @@ -334,6 +362,81 @@ TEST_CASE(MoeNvmeSchedulerFixture, split_model_reads_tensor_spans_from_multiple_ } #if !defined(_WIN32) +TEST_CASE(MoeNvmeSchedulerFixture, declared_shard_size_cannot_exceed_real_file) { + SyntheticModel model; + char path[] = "/tmp/moe_nvme_truncated_XXXXXX"; + const int fd = ::mkstemp(path); + NVME_REQUIRE(fd >= 0); + ::unlink(path); + std::vector bytes(4096, 0xa5); + NVME_REQUIRE(::write(fd, bytes.data(), bytes.size()) == (ssize_t) bytes.size()); + + MoeNvmeConfig config; + config.backend = MoeNvmeBackend::ThreadPool; + config.direct_io = MoeNvmeDirectMode::Disabled; + config.host_slots = 2; + MoeNvmeScheduler scheduler; + std::string err; + NVME_REQUIRE(scheduler.init(config, 32 * 1024, aligned_allocate, + aligned_free, nullptr, &err)); + NVME_REQUIRE(!scheduler.bind_source( + {nullptr, bytes.size() * 2, fd}, model.regions, &err)); + NVME_REQUIRE(err.find("shorter") != std::string::npos); + ::close(fd); +} + +#if defined(__linux__) && defined(O_DIRECT) +TEST_CASE(MoeNvmeSchedulerFixture, direct_io_accepts_valid_unaligned_shard_tail) { + constexpr size_t gate_bytes = 1000; + constexpr size_t up_bytes = 1000; + constexpr size_t down_bytes = 1300; + LayerExpertRegions layer; + layer.expert_bytes_gate = gate_bytes; + layer.expert_bytes_up = up_bytes; + layer.expert_bytes_down = down_bytes; + layer.gate_exps = {123, gate_bytes}; + layer.up_exps = {2125, up_bytes}; + layer.down_exps = {8192 + 37, down_bytes}; + std::vector file(layer.down_exps.offset + down_bytes, 0xa5); + fill_tensor(file, layer.gate_exps, gate_bytes, 0, 0, 1); + fill_tensor(file, layer.up_exps, up_bytes, 0, 1, 1); + fill_tensor(file, layer.down_exps, down_bytes, 0, 2, 1); + + char path[] = "/tmp/moe_nvme_direct_tail_XXXXXX"; + const int fd = ::mkstemp(path); + NVME_REQUIRE(fd >= 0); + ::unlink(path); + size_t written = 0; + while (written < file.size()) { + const ssize_t result = ::write( + fd, file.data() + written, file.size() - written); + NVME_REQUIRE(result > 0); + written += (size_t) result; + } + + MoeNvmeConfig config; + config.backend = MoeNvmeBackend::ThreadPool; + config.direct_io = MoeNvmeDirectMode::Enabled; + config.host_slots = 2; + config.io_threads = 1; + MoeNvmeScheduler scheduler; + std::string err; + NVME_REQUIRE(scheduler.init(config, gate_bytes + up_bytes + down_bytes, + aligned_allocate, aligned_free, nullptr, &err)); + NVME_REQUIRE(scheduler.bind_source( + {file.data(), file.size(), fd}, {layer}, &err)); + NVME_REQUIRE(scheduler.direct_io_active()); + MoeNvmeLease lease; + NVME_REQUIRE(scheduler.acquire(0, 0, lease, &err)); + verify_lease(lease, 0, 0); + lease.reset(); + NVME_REQUIRE(scheduler.stats().physical_bytes < 3 * 4096); + NVME_REQUIRE(scheduler.stats().errors == 0); + scheduler.destroy(); + ::close(fd); +} +#endif + TEST_CASE(MoeNvmeSchedulerFixture, split_real_files_use_the_fd_backend) { constexpr int experts = SyntheticModel::kExperts; std::vector shard_a(512 * 1024, 0xa5); @@ -374,7 +477,11 @@ TEST_CASE(MoeNvmeSchedulerFixture, split_real_files_use_the_fd_backend) { write_all(fd_b, shard_b); MoeNvmeConfig config; +#if defined(__linux__) + config.backend = MoeNvmeBackend::IoUring; +#else config.backend = MoeNvmeBackend::Auto; +#endif config.direct_io = MoeNvmeDirectMode::Disabled; config.host_slots = 4; config.io_threads = 2; @@ -423,6 +530,9 @@ TEST_CASE(MoeNvmeSchedulerFixture, real_file_backend_reads_exact_bytes) { aligned_free, nullptr, &err)); NVME_REQUIRE(scheduler.bind_source( {model.file.data(), model.file.size(), fd}, model.regions, &err)); +#if defined(__linux__) + NVME_REQUIRE(std::string(scheduler.effective_backend_name()) == "io_uring"); +#endif MoeNvmeLease lease; NVME_REQUIRE(scheduler.acquire(1, 7, lease, &err)); verify_lease(lease, 1, 7); From c8ad62a32ccd2d0a7b70b0edf3c9cd5a81e127d7 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sat, 1 Aug 2026 00:46:59 +0200 Subject: [PATCH 12/20] feat(server): add explicit MoE storage policy --- server/docs/DS4.md | 9 +- server/docs/ENVIRONMENT.md | 3 +- server/docs/MOE_NVME_STREAMING.md | 35 ++++-- server/src/common/backend_args.h | 8 ++ server/src/common/backend_factory.cpp | 44 ++++++- server/src/common/backend_factory.h | 3 + server/src/common/feature_gate.cpp | 17 +++ server/src/common/model_capabilities.h | 26 +++-- server/src/common/moe_storage_policy.h | 129 +++++++++++++++++++++ server/src/deepseek4/deepseek4_backend.cpp | 45 ++----- server/src/deepseek4/deepseek4_daemon.cpp | 13 +++ server/src/deepseek4/deepseek4_internal.h | 2 + server/src/kimi_k3/kimi_k3_backend.cpp | 6 +- server/src/kimi_k3/kimi_k3_backend.h | 5 +- server/src/server/server_main.cpp | 24 ++++ server/test/smoke_kimi_k3_forward.cpp | 5 +- server/test/test_feature_gate.cpp | 77 ++++++++++++ 17 files changed, 382 insertions(+), 69 deletions(-) create mode 100644 server/src/common/moe_storage_policy.h diff --git a/server/docs/DS4.md b/server/docs/DS4.md index 61fd5928b..32cb1cabf 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -180,9 +180,9 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. When the cold expert stack cannot fit on its compute device, the inference engine turns safe remaining memory into an adaptive warm-expert cache and streams only exact routed misses from NVMe. This supports both R9700+Strix -expert parallelism and a single Strix Halo. `DFLASH_MOE_NVME_COLD_TIER=auto` -selects streaming when capacity requires it; `on` forces at least one cold -expert per layer for qualification and `off` requires resident experts. On a +expert parallelism and a single Strix Halo. `--moe-storage auto` selects +streaming when capacity requires it; `ssd` forces at least one cold expert per +layer for qualification and `resident` prohibits SSD execution. On a full Lucebox the R9700 continues to own dense layers and hot experts. See [`MOE_NVME_STREAMING.md`](MOE_NVME_STREAMING.md) for the data path, tuning, and benchmark methodology. @@ -198,7 +198,8 @@ and benchmark methodology. | `DFLASH_DS4_MOE_TP` | Enable routed-expert partitioning. | | `DFLASH_DS4_MOE_TP_INPROC` | Use two local HIP backends instead of an expert IPC worker. | | `DFLASH_DS4_MOE_TP_GPU` | HIP device that owns the cold expert stack. | -| `DFLASH_MOE_NVME_COLD_TIER` | `auto`, `on`, or `off` for the SSD cold-capacity tier on dual-device or Strix-only deployments. | +| `DFLASH_MOE_STORAGE` | Environment equivalent of `--moe-storage auto|resident|ssd`; CLI takes precedence. | +| `DFLASH_MOE_NVME_COLD_TIER` | Deprecated compatibility alias (`auto`, `on`, `off`). | | `DFLASH_MOE_NVME_DEVICE_CACHE_MB` | Optional explicit adaptive device expert-cache budget; auto mode otherwise uses safe free memory. | | `DFLASH_EXPERT_BUDGET_MB` | Main-GPU memory budget for hot experts. | | `DFLASH_DS4_HOTNESS_CSV` | Optional per-layer routing profile for hot placement. | diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index e9672c20a..5a62c9145 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -32,7 +32,8 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_MOE_PLACEMENT` | unset | BURN-IN: offline placement JSON; listed experts belong to the primary owner. | | `DFLASH_MOE_PRIMARY_SHARE_PER_MILLE` | 500 | BURN-IN: deterministic primary route share when no explicit placement is supplied. | | `DFLASH_MOE_DUAL_STREAM_TRACE` | unset | DEBUG: per-layer dual-owner route counts and branch/wall timings. | -| `DFLASH_MOE_NVME_COLD_TIER` | auto | BURN-IN: DeepSeek4 cold-capacity policy (`auto`, `on`, `off`) for dual-device and Strix-only execution. | +| `DFLASH_MOE_STORAGE` | auto | Routed-MoE storage policy (`auto`, `resident`, `ssd`); prefer `--moe-storage`, which takes precedence. | +| `DFLASH_MOE_NVME_COLD_TIER` | unset | DEPRECATED: DeepSeek compatibility alias (`auto`, `on`, `off`). | | `DFLASH_MOE_NVME_*` | tuned defaults | BURN-IN: bounded MoE SSD scheduler/backend controls; see `MOE_NVME_STREAMING.md`. | | `GGML_CUDA_BATCH_PEER_COPIES` | unset | BURN-IN: publish ordered HIP peer copies with one cross-device dependency per source/destination pair. | | `DFLASH_MOE_PREFILL_PERSISTENT_OWNER_ALLOC` | 1 for qualified long heterogeneous prefill | KILL SWITCH: =0 restores per-layer route/owner scratch allocation. | diff --git a/server/docs/MOE_NVME_STREAMING.md b/server/docs/MOE_NVME_STREAMING.md index 47c2d50b0..229c0b01b 100644 --- a/server/docs/MOE_NVME_STREAMING.md +++ b/server/docs/MOE_NVME_STREAMING.md @@ -91,6 +91,25 @@ weight masks reconstruct the original route batch exactly. The native model router remains authoritative. Prediction may only issue a bounded prefetch; a wrong prediction cannot change model output. +## Storage policy + +The native server exposes one model-neutral operator policy: + +```bash +--moe-storage auto|resident|ssd +``` + +- `auto` preserves the model adapter's capacity-safe default and uses SSD when + the qualified adapter determines it is required. +- `resident` prohibits SSD expert execution and may fail allocation when the + routed weights do not fit available GPU/host memory. +- `ssd` forces the complete SSD-to-GPU capacity path. Startup fails before + model allocation when the architecture or placement lacks that path. + +CLI wins over `DFLASH_MOE_STORAGE`, which wins over the deprecated +`DFLASH_MOE_NVME_COLD_TIER` compatibility variable. Invalid values fail fast; +the resolved value and source are printed in the server configuration banner. + ## DeepSeek V4 Flash activation The existing heterogeneous mode is required. In `auto` mode, SSD streaming is @@ -102,14 +121,13 @@ export DFLASH_DS4_MOE_TP=1 export DFLASH_DS4_MOE_TP_INPROC=1 export DFLASH_DS4_MOE_TP_GPU=1 export DFLASH_EXPERT_BUDGET_MB=11700 -export DFLASH_MOE_NVME_COLD_TIER=auto ./build-hip-dual/dflash_server /path/to/model.gguf \ - --target-device hip:0 --peer-access + --target-device hip:0 --peer-access --moe-storage auto ``` -`DFLASH_MOE_NVME_COLD_TIER=on` forces the capacity tier for qualification; -`off` requires the old resident-cold path. In `auto`, a model that exceeds +`--moe-storage ssd` forces the capacity tier for qualification; +`resident` requires the old resident-cold path. In `auto`, a model that exceeds Strix receives all currently usable memory (after reserve) as its adaptive expert-cache budget, and only the remainder spills to SSD. @@ -120,13 +138,12 @@ exposes its GPU as `hip:0`: ```bash unset DFLASH_DS4_MOE_TP DFLASH_DS4_MOE_TP_INPROC DFLASH_DS4_MOE_TP_GPU -export DFLASH_MOE_NVME_COLD_TIER=on ./build-hip/dflash_server /path/to/model.gguf \ - --target-device hip:0 --max-ctx 8192 + --target-device hip:0 --max-ctx 8192 --moe-storage ssd ``` -`on` keeps at least one expert per layer in the SSD tier even if the model +`ssd` keeps at least one expert per layer in the SSD tier even if the model would otherwise be fully resident, making the path directly testable. For a model that genuinely exceeds UMA, `auto` chooses the partial placement itself. `DFLASH_EXPERT_BUDGET_MB` can additionally cap static routed-weight residency; @@ -143,7 +160,7 @@ export DFLASH_MOE_NVME_BACKEND=auto ./build-hip/dflash_server \ /path/to/Kimi-K3-UD-IQ1_S-00001-of-00014.gguf \ - --target-device hip:0 --max-ctx 8192 + --target-device hip:0 --max-ctx 8192 --moe-storage auto ``` No DeepSeek MoE-TP variables are required. Cache memory is chosen only after @@ -186,6 +203,8 @@ not improve the qualified P310 drive and consumes extra pinned/system memory. | Variable | Default | Meaning | |---|---:|---| +| `DFLASH_MOE_STORAGE` | `auto` | Environment equivalent of `--moe-storage`; CLI takes precedence | +| `DFLASH_MOE_NVME_COLD_TIER` | unset | Deprecated DeepSeek compatibility alias (`auto`, `on`, `off`) | | `DFLASH_MOE_NVME_BACKEND` | `auto` | `auto`, `uring`, `pread`, or `mmap` | | `DFLASH_MOE_NVME_DIRECT` | `auto` | `auto`, `on`, or `off` | | `DFLASH_MOE_NVME_SLOTS` | `8` | Fixed pinned host slots | diff --git a/server/src/common/backend_args.h b/server/src/common/backend_args.h index 1607b1495..c0573d2de 100644 --- a/server/src/common/backend_args.h +++ b/server/src/common/backend_args.h @@ -9,6 +9,9 @@ #include "placement/remote_draft_config.h" #include "placement/remote_target_shard_config.h" #include "prefill_attention_mode.h" +#include "moe_storage_policy.h" + +#include namespace dflash::common { @@ -55,6 +58,11 @@ struct BackendArgs { int ds4_expert_top_k = 0; // 0 = model default bool ds4_fused_decode = false; + // Routed-MoE capacity policy. An unset value allows the factory to apply + // environment compatibility and the default; an explicit Auto must remain + // distinguishable because CLI always wins over environment configuration. + std::optional moe_storage; + // Attention and speculative-decode options. Individual backends consume // only the fields they support. int fa_window = 0; // 0 = full attention. qwen3.6 full-attn layers must see the whole context; a finite window drops the system prompt/tools -> breaks tool calls. diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 15583e06f..a24fa88ef 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -18,6 +18,7 @@ #include "qwen35_layer_split_adapter.h" #include +#include #include #include @@ -49,6 +50,7 @@ DFLASH_ARCH_FIELD_TRAIT(has_verify_width, verify_width); DFLASH_ARCH_FIELD_TRAIT(has_draft_swa, draft_swa_window); DFLASH_ARCH_FIELD_TRAIT(has_ddtree_mode, ddtree_mode); DFLASH_ARCH_FIELD_TRAIT(has_max_verify_tokens, max_verify_tokens); +DFLASH_ARCH_FIELD_TRAIT(has_moe_storage, moe_storage); #undef DFLASH_ARCH_FIELD_TRAIT @@ -87,7 +89,8 @@ constexpr bool layer_split_carries(FeatureSupport support) { DFLASH_CHECK_ARCH_OPTION(arch_name, Mono, Split, has_ddtree, ddtree); \ DFLASH_CHECK_ARCH_OPTION(arch_name, Mono, Split, has_verify_width, verify_width); \ DFLASH_CHECK_ARCH_OPTION(arch_name, Mono, Split, has_fa_window, fa_window); \ - DFLASH_CHECK_ARCH_OPTION(arch_name, Mono, Split, has_draft_swa, draft_swa) + DFLASH_CHECK_ARCH_OPTION(arch_name, Mono, Split, has_draft_swa, draft_swa); \ + DFLASH_CHECK_ARCH_OPTION(arch_name, Mono, Split, has_moe_storage, moe_ssd_storage) DFLASH_CHECK_ARCH("qwen35", Qwen35Config, Qwen35LayerSplitAdapterConfig); DFLASH_CHECK_ARCH("qwen35moe", Qwen35Config, NoLayerSplitConfig); @@ -108,6 +111,14 @@ PlacementBackend resolve_target_backend( : args.device.backend; } +MoeStoragePolicyResolution resolve_requested_moe_storage( + const BackendArgs & args) { + return resolve_moe_storage_policy( + args.moe_storage, + std::getenv(kMoeStorageEnvironment), + std::getenv(kLegacyMoeStorageEnvironment)); +} + } // namespace std::string detect_arch(const char * model_path) { @@ -131,6 +142,12 @@ BackendPreparation prepare_backend( preparation.plan.target_backend_ = resolve_target_backend( args, preparation.plan.compiled_backend_); preparation.plan.model_ = inspect_gguf_model_info(args.model_path); + preparation.plan.moe_storage_ = resolve_requested_moe_storage(args); + if (!preparation.plan.moe_storage_.ok()) { + preparation.error = BackendPreparationError::InvalidRequest; + preparation.message = preparation.plan.moe_storage_.error; + return preparation; + } if (preparation.plan.arch().empty()) { preparation.error = BackendPreparationError::ModelInspection; @@ -152,8 +169,10 @@ BackendPreparation prepare_backend( return preparation; } + BackendArgs effective_args = args; + effective_args.moe_storage = preparation.plan.moe_storage_policy(); preparation.message = check_feature_compatibility( - args, + effective_args, preparation.plan.features(), preparation.plan.arch(), preparation.plan.target_backend(), @@ -164,7 +183,12 @@ BackendPreparation prepare_backend( } preparation.warnings = collect_feature_warnings( - args, preparation.plan.features(), preparation.plan.arch()); + effective_args, preparation.plan.features(), preparation.plan.arch()); + if (!preparation.plan.moe_storage_.warning.empty()) { + preparation.warnings.insert( + preparation.warnings.begin(), + preparation.plan.moe_storage_.warning); + } return preparation; } @@ -202,6 +226,14 @@ std::unique_ptr create_backend( "[backend_factory] resolved plan does not match target placement\n"); return nullptr; } + const MoeStoragePolicyResolution current_storage = + resolve_requested_moe_storage(args); + if (!current_storage.ok() || + current_storage.policy != plan.moe_storage_policy()) { + std::fprintf(stderr, + "[backend_factory] resolved plan does not match MoE storage policy\n"); + return nullptr; + } const std::string & arch = plan.arch(); if (arch.empty()) { @@ -215,8 +247,10 @@ std::unique_ptr create_backend( // Recheck at the construction boundary in case raw arguments changed // after preparation. No entry point can dispatch an incoherent request. + BackendArgs effective_args = args; + effective_args.moe_storage = plan.moe_storage_policy(); const std::string incompatible = check_feature_compatibility( - args, + effective_args, plan.features(), arch, plan.target_backend(), @@ -414,6 +448,7 @@ std::unique_ptr create_backend( cfg.expert_top_k = args.ds4_expert_top_k; cfg.fused_decode = args.ds4_fused_decode; cfg.prefill_mode = args.ds4_prefill_mode; + cfg.moe_storage = plan.moe_storage_policy(); auto backend = std::make_unique(cfg); if (!backend->init()) { @@ -443,6 +478,7 @@ std::unique_ptr create_backend( cfg.model_path = args.model_path; cfg.device = args.device; cfg.stream_fd = args.stream_fd; + cfg.moe_storage = plan.moe_storage_policy(); auto backend = std::make_unique(cfg); if (!backend->init()) { diff --git a/server/src/common/backend_factory.h b/server/src/common/backend_factory.h index bb3d71164..b5c21e0cf 100644 --- a/server/src/common/backend_factory.h +++ b/server/src/common/backend_factory.h @@ -34,11 +34,14 @@ class ResolvedBackendPlan { PlacementBackend target_backend() const { return target_backend_; } PlacementBackend compiled_backend() const { return compiled_backend_; } const BackendFeatureConfig & features() const { return features_; } + MoeStoragePolicy moe_storage_policy() const { return moe_storage_.policy; } + MoeStoragePolicySource moe_storage_source() const { return moe_storage_.source; } private: std::string model_path_; GgufModelInfo model_; BackendFeatureConfig features_; + MoeStoragePolicyResolution moe_storage_; PlacementBackend target_backend_ = PlacementBackend::Auto; PlacementBackend compiled_backend_ = PlacementBackend::Auto; diff --git a/server/src/common/feature_gate.cpp b/server/src/common/feature_gate.cpp index 6ab697b45..c85e7b5be 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -121,6 +121,23 @@ std::string check_feature_compatibility( placement_device_name(args.device) + " alone"; } + // ── SSD-backed routed experts × architecture/dispatch path + // Auto and resident are safe everywhere. Forcing SSD must reach a complete + // storage+compute adapter; a prefetch-only integration is not sufficient. + const MoeStoragePolicy storage = + args.moe_storage.value_or(MoeStoragePolicy::Auto); + const bool split_dispatch = + args.device.is_layer_split() || args.remote_target_shard.enabled(); + if (storage == MoeStoragePolicy::Ssd && + !arch_supports_moe_ssd_storage(arch, split_dispatch)) { + if (split_dispatch && arch_supports_moe_ssd_storage(arch, false)) { + return "--moe-storage ssd is supported for architecture '" + arch + + "' only on monolithic placement"; + } + return "model architecture '" + arch + + "' does not support --moe-storage ssd"; + } + // ── remote draft execution × architecture if (args.remote_draft.enabled() && args.draft_path && !arch_supports_remote_draft(arch)) { diff --git a/server/src/common/model_capabilities.h b/server/src/common/model_capabilities.h index bc956e7aa..7187e9921 100644 --- a/server/src/common/model_capabilities.h +++ b/server/src/common/model_capabilities.h @@ -55,6 +55,7 @@ struct ArchCapabilities { FeatureSupport verify_width; // --verify-width FeatureSupport fa_window; // --fa-window FeatureSupport draft_swa; // --draft-swa + FeatureSupport moe_ssd_storage; // --moe-storage ssd }; inline constexpr FeatureSupport kNever = FeatureSupport::Never; @@ -62,14 +63,14 @@ inline constexpr FeatureSupport kMono = FeatureSupport::Monolithic; inline constexpr FeatureSupport kBoth = FeatureSupport::Both; inline constexpr ArchCapabilities kArchCapabilities[] = { -// arch split rdraft pflash offload draft ddtree vwidth fa_win dswa - {"qwen35", true, true, true, false, kBoth, kBoth, kNever, kBoth, kBoth}, - {"qwen35moe", false, false, false, true, kMono, kMono, kNever, kMono, kMono}, - {"laguna", true, false, false, true, kMono, kMono, kMono, kNever, kNever}, - {"qwen3", false, false, true, false, kNever, kNever, kNever, kNever, kNever}, - {"gemma4", true, false, false, false, kMono, kNever, kNever, kBoth, kNever}, - {"deepseek4", true, false, false, false, kNever, kNever, kNever, kNever, kNever}, - {"kimi-k3", false, false, false, false, kNever, kNever, kNever, kNever, kNever}, +// arch split rdraft pflash offload draft ddtree vwidth fa_win dswa moe-ssd + {"qwen35", true, true, true, false, kBoth, kBoth, kNever, kBoth, kBoth, kNever}, + {"qwen35moe", false, false, false, true, kMono, kMono, kNever, kMono, kMono, kNever}, + {"laguna", true, false, false, true, kMono, kMono, kMono, kNever, kNever, kNever}, + {"qwen3", false, false, true, false, kNever, kNever, kNever, kNever, kNever, kNever}, + {"gemma4", true, false, false, false, kMono, kNever, kNever, kBoth, kNever, kNever}, + {"deepseek4", true, false, false, false, kNever, kNever, kNever, kNever, kNever, kMono}, + {"kimi-k3", false, false, false, false, kNever, kNever, kNever, kNever, kNever, kMono}, }; inline constexpr std::size_t kArchCount = @@ -107,7 +108,8 @@ constexpr bool row_has_both(const ArchCapabilities & c) { c.ddtree == FeatureSupport::Both || c.verify_width == FeatureSupport::Both || c.fa_window == FeatureSupport::Both || - c.draft_swa == FeatureSupport::Both; + c.draft_swa == FeatureSupport::Both || + c.moe_ssd_storage == FeatureSupport::Both; } constexpr bool table_rows_named() { @@ -236,4 +238,10 @@ inline bool arch_supports_draft_swa(const std::string & arch, return detail::arch_has(arch, &ArchCapabilities::draft_swa, is_layer_split); } +inline bool arch_supports_moe_ssd_storage(const std::string & arch, + bool split_dispatch) { + return detail::arch_has( + arch, &ArchCapabilities::moe_ssd_storage, split_dispatch); +} + } // namespace dflash::common diff --git a/server/src/common/moe_storage_policy.h b/server/src/common/moe_storage_policy.h new file mode 100644 index 000000000..f013963fe --- /dev/null +++ b/server/src/common/moe_storage_policy.h @@ -0,0 +1,129 @@ +// Operator policy for routed-MoE expert storage. +// +// Parsing and precedence live here so model backends consume one resolved, +// typed value. They must not inspect CLI or environment strings themselves. + +#pragma once + +#include +#include +#include + +namespace dflash::common { + +inline constexpr const char * kMoeStorageEnvironment = + "DFLASH_MOE_STORAGE"; +inline constexpr const char * kLegacyMoeStorageEnvironment = + "DFLASH_MOE_NVME_COLD_TIER"; + +enum class MoeStoragePolicy { + Auto, + Resident, + Ssd, +}; + +inline const char * moe_storage_policy_name(MoeStoragePolicy policy) { + switch (policy) { + case MoeStoragePolicy::Auto: return "auto"; + case MoeStoragePolicy::Resident: return "resident"; + case MoeStoragePolicy::Ssd: return "ssd"; + } + return "unknown"; +} + +inline bool parse_moe_storage_policy(std::string_view value, + MoeStoragePolicy & out) { + if (value == "auto") { + out = MoeStoragePolicy::Auto; + return true; + } + if (value == "resident") { + out = MoeStoragePolicy::Resident; + return true; + } + if (value == "ssd") { + out = MoeStoragePolicy::Ssd; + return true; + } + return false; +} + +enum class MoeStoragePolicySource { + Default, + LegacyEnvironment, + Environment, + Cli, +}; + +inline const char * moe_storage_policy_source_name( + MoeStoragePolicySource source) { + switch (source) { + case MoeStoragePolicySource::Default: return "default"; + case MoeStoragePolicySource::LegacyEnvironment: return "legacy environment"; + case MoeStoragePolicySource::Environment: return "environment"; + case MoeStoragePolicySource::Cli: return "CLI"; + } + return "unknown"; +} + +struct MoeStoragePolicyResolution { + MoeStoragePolicy policy = MoeStoragePolicy::Auto; + MoeStoragePolicySource source = MoeStoragePolicySource::Default; + std::string error; + std::string warning; + + bool ok() const { return error.empty(); } +}; + +// Resolution order is intentionally explicit and testable: +// CLI > DFLASH_MOE_STORAGE > legacy DFLASH_MOE_NVME_COLD_TIER > auto. +// Environment values are parameters rather than read internally, keeping the +// function deterministic and avoiding hidden process-global state in tests. +inline MoeStoragePolicyResolution resolve_moe_storage_policy( + std::optional cli, + const char * environment, + const char * legacy_environment) { + MoeStoragePolicyResolution out; + if (cli.has_value()) { + out.policy = *cli; + out.source = MoeStoragePolicySource::Cli; + return out; + } + + if (environment && environment[0] != '\0') { + if (!parse_moe_storage_policy(environment, out.policy)) { + out.error = "DFLASH_MOE_STORAGE expects auto, resident, or ssd; got '" + + std::string(environment) + "'"; + return out; + } + out.source = MoeStoragePolicySource::Environment; + return out; + } + + if (!legacy_environment || legacy_environment[0] == '\0') { + return out; + } + + const std::string_view legacy(legacy_environment); + if (legacy == "auto") { + out.policy = MoeStoragePolicy::Auto; + } else if (legacy == "1" || legacy == "on" || legacy == "true" || + legacy == "ssd") { + out.policy = MoeStoragePolicy::Ssd; + } else if (legacy == "0" || legacy == "off" || legacy == "false" || + legacy == "resident") { + out.policy = MoeStoragePolicy::Resident; + } else { + out.error = + "DFLASH_MOE_NVME_COLD_TIER expects auto, on, or off; got '" + + std::string(legacy_environment) + "'"; + return out; + } + out.source = MoeStoragePolicySource::LegacyEnvironment; + out.warning = + "DFLASH_MOE_NVME_COLD_TIER is deprecated; use --moe-storage or " + "DFLASH_MOE_STORAGE"; + return out; +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index ea4d831c1..48cc79853 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -40,29 +40,6 @@ static bool env_flag_enabled(const char * name) { return value && value[0] && std::strcmp(value, "0") != 0; } -enum class MoeNvmeColdTierMode { - Auto, - Enabled, - Disabled, - Invalid, -}; - -static MoeNvmeColdTierMode moe_nvme_cold_tier_mode() { - const char * value = std::getenv("DFLASH_MOE_NVME_COLD_TIER"); - if (!value || !value[0] || std::strcmp(value, "auto") == 0) { - return MoeNvmeColdTierMode::Auto; - } - if (std::strcmp(value, "1") == 0 || std::strcmp(value, "on") == 0 || - std::strcmp(value, "true") == 0) { - return MoeNvmeColdTierMode::Enabled; - } - if (std::strcmp(value, "0") == 0 || std::strcmp(value, "off") == 0 || - std::strcmp(value, "false") == 0) { - return MoeNvmeColdTierMode::Disabled; - } - return MoeNvmeColdTierMode::Invalid; -} - static void configure_gfx1151_dspark_mmvq_default(int gpu) { #if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) if (!env_flag_enabled("DFLASH_DS4_SPEC") || @@ -549,7 +526,7 @@ bool DeepSeek4Backend::load_model() { const bool force_full = env_flag_enabled("DFLASH_DS4_FORCE_FULL_LOAD"); const bool heterogeneous_tp = env_flag_enabled("DFLASH_DS4_MOE_TP"); const bool explicit_ssd_capacity = - moe_nvme_cold_tier_mode() == MoeNvmeColdTierMode::Enabled; + cfg_.moe_storage == MoeStoragePolicy::Ssd; const bool need_monolithic = requires_monolithic_model() && !heterogeneous_tp && !explicit_ssd_capacity; if (target_backend == PlacementBackend::Hip && @@ -868,9 +845,9 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & const bool all_cold = env_flag_enabled("DFLASH_DS4_MOE_TP_ALL_COLD"); int hot_per_layer = all_cold ? 0 : budget.max_hot_per_layer; if (!all_cold && - moe_nvme_cold_tier_mode() == MoeNvmeColdTierMode::Enabled && + cfg_.moe_storage == MoeStoragePolicy::Ssd && hot_per_layer >= w.n_expert) { - // `on` is also a qualification switch. Keep one exact expert per + // `ssd` is also a qualification switch. Keep one exact expert per // layer in the SSD tier when the whole model would otherwise fit, so // a single-Strix user can exercise the real path without inventing a // fragile machine-specific memory cap. `auto` still prefers full @@ -951,19 +928,13 @@ bool DeepSeek4Backend::init_hybrid_model() { const bool inprocess_tp = tp_requested && ds4_inprocess_moe_tp_enabled(); const bool external_tp = tp_requested && !inprocess_tp; - MoeNvmeColdTierMode nvme_mode = moe_nvme_cold_tier_mode(); - if (nvme_mode == MoeNvmeColdTierMode::Invalid) { - std::fprintf(stderr, - "[deepseek4] ignoring invalid DFLASH_MOE_NVME_COLD_TIER; using auto\n"); - nvme_mode = MoeNvmeColdTierMode::Auto; - } - // A partially resident single-GPU model streams its non-resident experts // on the same device. In-process TP instead streams them on the expert GPU; // external TP leaves them to the remote worker and does not start a local - // SSD service. `off` explicitly selects the older materialized CPU tail. + // SSD service. `resident` explicitly selects the older materialized CPU + // tail. bool stream_cold = !external_tp && - nvme_mode != MoeNvmeColdTierMode::Disabled; + cfg_.moe_storage != MoeStoragePolicy::Resident; if (inprocess_tp) { const int expert_gpu = ds4_moe_tp_gpu(cfg_.device.gpu); if (expert_gpu < 0) return false; @@ -989,8 +960,8 @@ bool DeepSeek4Backend::init_hybrid_model() { // On a separate expert GPU, auto mode retains the established fully // resident path whenever the cold stack fits after a conservative - // reserve. Explicit on/off remain authoritative. - if (nvme_mode == MoeNvmeColdTierMode::Auto) { + // reserve. Explicit ssd/resident modes remain authoritative. + if (cfg_.moe_storage == MoeStoragePolicy::Auto) { Ds4ExpertMemoryInfo info; std::string memory_error; size_t expert_free = 0; diff --git a/server/src/deepseek4/deepseek4_daemon.cpp b/server/src/deepseek4/deepseek4_daemon.cpp index fabc1c184..a86f43540 100644 --- a/server/src/deepseek4/deepseek4_daemon.cpp +++ b/server/src/deepseek4/deepseek4_daemon.cpp @@ -5,6 +5,7 @@ #include "common/daemon_loop.h" #include +#include namespace dflash::common { @@ -19,6 +20,18 @@ int run_deepseek4_daemon(const char * model_path, cfg.stream_fd = stream_fd; cfg.max_ctx = max_ctx; cfg.chunk = chunk > 0 ? chunk : 512; + const MoeStoragePolicyResolution storage = resolve_moe_storage_policy( + {}, std::getenv(kMoeStorageEnvironment), + std::getenv(kLegacyMoeStorageEnvironment)); + if (!storage.ok()) { + std::fprintf(stderr, "[deepseek4-daemon] %s\n", storage.error.c_str()); + return 2; + } + if (!storage.warning.empty()) { + std::fprintf(stderr, "[deepseek4-daemon] warning: %s\n", + storage.warning.c_str()); + } + cfg.moe_storage = storage.policy; auto backend = std::make_unique(cfg); if (!backend->init()) { diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 0e9beac0b..a3344c4a3 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -24,6 +24,7 @@ #include "internal.h" #include "common/layer_split_utils.h" +#include "common/moe_storage_policy.h" #include "common/prefill_attention_mode.h" namespace dflash::common { @@ -309,6 +310,7 @@ struct DeepSeek4BackendConfig { int max_ctx = 0; // 0 = auto from SWA + compression capacity int expert_top_k = 0; // 0 = use all model-routed experts bool fused_decode = false; // single-graph GPU decode + MoeStoragePolicy moe_storage = MoeStoragePolicy::Auto; }; // ─── Function declarations ────────────────────────────────────────────── diff --git a/server/src/kimi_k3/kimi_k3_backend.cpp b/server/src/kimi_k3/kimi_k3_backend.cpp index 3e90eedff..6395086b2 100644 --- a/server/src/kimi_k3/kimi_k3_backend.cpp +++ b/server/src/kimi_k3/kimi_k3_backend.cpp @@ -304,9 +304,11 @@ bool KimiK3Backend::init() { cfg_.device.primary_gpu()); return false; } + const bool stream_routed_experts = + cfg_.moe_storage != MoeStoragePolicy::Resident; if (!load_kimi_k3_gguf( cfg_.model_path, backend_, weights_, - cfg_.stream_routed_experts)) { + stream_routed_experts)) { std::fprintf(stderr, "[kimi-k3] model load failed: %s\n", dflash27b_last_error()); return false; @@ -354,7 +356,7 @@ bool KimiK3Backend::unpark(ParkTarget target) { if (parked_) { if (!load_kimi_k3_gguf( cfg_.model_path, backend_, weights_, - cfg_.stream_routed_experts) || + cfg_.moe_storage != MoeStoragePolicy::Resident) || (weights_.routed_experts_streamed && !init_streaming())) { return false; } diff --git a/server/src/kimi_k3/kimi_k3_backend.h b/server/src/kimi_k3/kimi_k3_backend.h index 2026f7b2e..64e7453ef 100644 --- a/server/src/kimi_k3/kimi_k3_backend.h +++ b/server/src/kimi_k3/kimi_k3_backend.h @@ -2,6 +2,7 @@ #include "common/model_backend.h" #include "common/moe_hybrid_stream.h" +#include "common/moe_storage_policy.h" #include "kimi_k3_internal.h" #include "placement/placement_config.h" @@ -19,9 +20,9 @@ struct KimiK3BackendConfig { // routed work is partitioned between both GPUs while dense KDA/MLA, // recurrent state, and sampling remain primary-owned. int expert_gpu = -1; - // Production Kimi uses file-backed routed experts. The resident mode is + // Auto uses Kimi's capacity-safe file-backed routed experts. Resident is // retained as a deterministic oracle for small architecture fixtures. - bool stream_routed_experts = true; + MoeStoragePolicy moe_storage = MoeStoragePolicy::Auto; }; class KimiK3Backend final : public ModelBackend { diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 6d654848d..911c9243f 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -100,6 +100,8 @@ static void print_usage(const char * prog) { " --ds4-prefill DeepSeek4 prefill: exact, dense, or sparse\n" " (default: exact; dense/sparse are experimental\n" " and may change generated tokens)\n" + " --moe-storage Routed-MoE storage: auto, resident, or ssd\n" + " (default: auto; env: DFLASH_MOE_STORAGE)\n" " --fa-window Flash-attention sliding window (default: 0=full).\n" " WARNING: >0 drops system prompt / tool definitions\n" " from attention at long contexts. Use 0 for tools.\n" @@ -318,6 +320,24 @@ int main(int argc, char ** argv) { bargs.device.peer_access = true; } else if (std::strcmp(argv[i], "--chunk") == 0 && i + 1 < argc) { bargs.chunk = std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--moe-storage") == 0) { + MoeStoragePolicy policy; + if (i + 1 >= argc || + !parse_moe_storage_policy(argv[i + 1], policy)) { + std::fprintf(stderr, + "[server] --moe-storage expects auto, resident, or ssd\n"); + return 2; + } + ++i; + bargs.moe_storage = policy; + } else if (std::strncmp(argv[i], "--moe-storage=", 14) == 0) { + MoeStoragePolicy policy; + if (!parse_moe_storage_policy(argv[i] + 14, policy)) { + std::fprintf(stderr, + "[server] --moe-storage expects auto, resident, or ssd\n"); + return 2; + } + bargs.moe_storage = policy; } else if (std::strcmp(argv[i], "--ds4-fused-decode") == 0) { bargs.ds4_fused_decode = true; } else if (std::strcmp(argv[i], "--ds4-expert-top-k") == 0 && i + 1 < argc) { @@ -968,6 +988,10 @@ int main(int argc, char ** argv) { sconfig.effort_tiers.max, src_of(cli_set.effort_max)); std::fprintf(stderr, "[server] │ target_device = %s\n", placement_device_name(bargs.device).c_str()); + std::fprintf(stderr, "[server] │ moe_storage = %s (%s)\n", + moe_storage_policy_name(backend_plan.moe_storage_policy()), + moe_storage_policy_source_name( + backend_plan.moe_storage_source())); if (bargs.device.is_layer_split()) { std::fprintf(stderr, "[server] │ target_shards ="); for (size_t i = 0; i < bargs.device.layer_split_gpus.size(); ++i) { diff --git a/server/test/smoke_kimi_k3_forward.cpp b/server/test/smoke_kimi_k3_forward.cpp index e8d258768..c635a84c8 100644 --- a/server/test/smoke_kimi_k3_forward.cpp +++ b/server/test/smoke_kimi_k3_forward.cpp @@ -37,8 +37,9 @@ int main(int argc, char ** argv) { config.model_path = model; config.device.gpu = gpu; config.device.max_ctx = 4096; - config.stream_routed_experts = - argc <= 5 || std::atoi(argv[5]) != 0; + config.moe_storage = argc <= 5 || std::atoi(argv[5]) != 0 + ? MoeStoragePolicy::Ssd + : MoeStoragePolicy::Resident; config.expert_gpu = argc > 6 ? std::atoi(argv[6]) : -1; KimiK3Backend backend(config); if (!backend.init()) return 1; diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp index 67daef4fe..3f7c2c560 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -43,6 +43,43 @@ static int test_count = 0; // One case per rule cluster in check_feature_compatibility(). All resolved // facts are parameters, so none of this needs a model file or GPU. +static void test_moe_storage_policy_resolution() { + MoeStoragePolicy parsed = MoeStoragePolicy::Auto; + TEST_ASSERT(parse_moe_storage_policy("auto", parsed)); + TEST_ASSERT(parsed == MoeStoragePolicy::Auto); + TEST_ASSERT(parse_moe_storage_policy("resident", parsed)); + TEST_ASSERT(parsed == MoeStoragePolicy::Resident); + TEST_ASSERT(parse_moe_storage_policy("ssd", parsed)); + TEST_ASSERT(parsed == MoeStoragePolicy::Ssd); + TEST_ASSERT(!parse_moe_storage_policy("disk", parsed)); + + MoeStoragePolicyResolution resolved = + resolve_moe_storage_policy({}, nullptr, nullptr); + TEST_ASSERT(resolved.ok()); + TEST_ASSERT(resolved.policy == MoeStoragePolicy::Auto); + TEST_ASSERT(resolved.source == MoeStoragePolicySource::Default); + + resolved = resolve_moe_storage_policy({}, "resident", "on"); + TEST_ASSERT(resolved.ok()); + TEST_ASSERT(resolved.policy == MoeStoragePolicy::Resident); + TEST_ASSERT(resolved.source == MoeStoragePolicySource::Environment); + + resolved = resolve_moe_storage_policy( + MoeStoragePolicy::Ssd, "invalid", "off"); + TEST_ASSERT(resolved.ok()); + TEST_ASSERT(resolved.policy == MoeStoragePolicy::Ssd); + TEST_ASSERT(resolved.source == MoeStoragePolicySource::Cli); + + resolved = resolve_moe_storage_policy({}, nullptr, "on"); + TEST_ASSERT(resolved.ok()); + TEST_ASSERT(resolved.policy == MoeStoragePolicy::Ssd); + TEST_ASSERT(resolved.source == MoeStoragePolicySource::LegacyEnvironment); + TEST_ASSERT(!resolved.warning.empty()); + + TEST_ASSERT(!resolve_moe_storage_policy({}, "invalid", nullptr).ok()); + TEST_ASSERT(!resolve_moe_storage_policy({}, nullptr, "invalid").ok()); +} + static BackendArgs gate_args_hip_deepseek4() { BackendArgs args; args.model_path = "/nonexistent/model.gguf"; @@ -279,6 +316,38 @@ static void test_feature_gate_layer_split_requires_supported_arch() { TEST_ASSERT(gate_result(single, "kimi-k3", PlacementBackend::Cuda).empty()); } +static void test_feature_gate_moe_ssd_requires_complete_adapter() { + BackendArgs args; + args.model_path = "/nonexistent/model.gguf"; + args.moe_storage = MoeStoragePolicy::Ssd; + + TEST_ASSERT(gate_result( + args, "deepseek4", PlacementBackend::Hip).empty()); + TEST_ASSERT(gate_result( + args, "kimi-k3", PlacementBackend::Hip).empty()); + for (const char * arch : {"qwen35", "qwen35moe", "laguna", + "qwen3", "gemma4"}) { + TEST_ASSERT(!gate_result(args, arch, PlacementBackend::Hip).empty()); + } + + BackendArgs split = args; + TEST_ASSERT(parse_placement_device_list("hip:0,hip:1", split.device)); + TEST_ASSERT(!gate_result( + split, "deepseek4", PlacementBackend::Hip).empty()); + + BackendArgs remote = args; + remote.remote_target_shard.ipc_bin = "/usr/bin/target-shard"; + TEST_ASSERT(!gate_result( + remote, "deepseek4", PlacementBackend::Hip).empty()); + + // Resident prohibits SSD streaming and is therefore safe everywhere. + args.moe_storage = MoeStoragePolicy::Resident; + TEST_ASSERT(gate_result( + args, "qwen35", PlacementBackend::Hip).empty()); + TEST_ASSERT(gate_result( + args, "laguna", PlacementBackend::Hip).empty()); +} + // ── Inert-flag warnings ───────────────────────────────────────────────── // Warnings must never gate admission, so each case also asserts the same // configuration passes check_feature_compatibility(). @@ -406,10 +475,17 @@ static void test_model_capability_tables() { TEST_ASSERT(!arch_supports_verify_width("qwen36", false)); TEST_ASSERT(!arch_supports_fa_window("qwen36", false)); TEST_ASSERT(!arch_supports_draft_swa("qwen36", false)); + TEST_ASSERT(!arch_supports_moe_ssd_storage("qwen36", false)); + + TEST_ASSERT(arch_supports_moe_ssd_storage("deepseek4", false)); + TEST_ASSERT(!arch_supports_moe_ssd_storage("deepseek4", true)); + TEST_ASSERT(arch_supports_moe_ssd_storage("kimi-k3", false)); + TEST_ASSERT(!arch_supports_moe_ssd_storage("qwen35moe", false)); } int main() { std::fprintf(stderr, "\n\u2500\u2500 Backend feature/architecture gate \u2500\u2500\n"); + RUN_TEST(test_moe_storage_policy_resolution); RUN_TEST(test_feature_gate_accepts_plain_launch); RUN_TEST(test_feature_gate_rejects_undetected_arch); RUN_TEST(test_feature_gate_requires_compiled_target_backend); @@ -422,6 +498,7 @@ int main() { RUN_TEST(test_feature_gate_ds4_decode_options_require_monolithic_hip); RUN_TEST(test_feature_gate_remote_draft_requires_supported_arch); RUN_TEST(test_feature_gate_layer_split_requires_supported_arch); + RUN_TEST(test_feature_gate_moe_ssd_requires_complete_adapter); RUN_TEST(test_feature_warnings_silent_when_supported); RUN_TEST(test_feature_warnings_report_inert_draft); RUN_TEST(test_feature_warnings_report_inert_decode_tunables); From 9abec86440a8d1a86ef60ef7235ed5f0396ba8e5 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sat, 1 Aug 2026 00:58:16 +0200 Subject: [PATCH 13/20] bench(kimi): add SSD deployment comparison harness --- server/docs/KIMI_K3_HETERO.md | 33 +- .../scripts/benchmark_kimi_k3_deployments.py | 573 ++++++++++++++++++ .../test_benchmark_kimi_k3_deployments.py | 109 ++++ 3 files changed, 711 insertions(+), 4 deletions(-) create mode 100644 server/scripts/benchmark_kimi_k3_deployments.py create mode 100644 server/scripts/test_benchmark_kimi_k3_deployments.py diff --git a/server/docs/KIMI_K3_HETERO.md b/server/docs/KIMI_K3_HETERO.md index 1a178bee7..e69a9097f 100644 --- a/server/docs/KIMI_K3_HETERO.md +++ b/server/docs/KIMI_K3_HETERO.md @@ -176,7 +176,32 @@ router, or per-layer worker creation. Full-scale qualification requires: 4. Replace the correctness-first host partial join with a device-resident peer join, then tune the placement until the two owner branches balance. -The full 594 GB model cannot currently be staged on the qualification box, -which currently has about 513 GB free. It needs at least about 650 GB of safe -free space for all shards plus logging/headroom; existing user models should -not be deleted implicitly. +Staging the released quant requires 594 GB for the shards plus operational +headroom. The benchmark validates that every shard is present before starting; +it never downloads models or deletes existing data implicitly. + +## Reproducible full-model comparison + +Once all 14 shards are present, one harness runs the two relevant deployments +serially so they cannot contend for the SSD or GPUs: + +```bash +python3 server/scripts/benchmark_kimi_k3_deployments.py \ + /models/Kimi-K3-UD-IQ1_S/Kimi-K3-UD-IQ1_S-00001-of-00014.gguf \ + --server-bin server/build-hip-dual/dflash_server \ + --r9700-device 0 --strix-device 1 \ + --output-dir bench-out/kimi-k3-lucebox +``` + +The profiles are Strix-only+SSD and heterogeneous Strix+R9700+SSD. For the +published IQ1_S checkpoint, Strix is the capacity-safe primary because the +non-routed tensors alone exceed the R9700's memory; R9700 concurrently owns a +partition of routed experts. `--hetero-primary r9700` is available for a later +checkpoint whose non-routed plan fits the discrete GPU. + +Each profile starts from a fresh server, forces `--moe-storage ssd`, clears +inherited MoE tuning, disables HTTP/prefix caches, and issues one cold followed +by two warm deterministic requests. The output directory contains the command, +effective MoE environment, client first-event time, server prefill/decode rate, +complete server log, and parsed per-owner NVMe/cache/graph telemetry. A +deterministic output mismatch between profiles fails the benchmark. diff --git a/server/scripts/benchmark_kimi_k3_deployments.py b/server/scripts/benchmark_kimi_k3_deployments.py new file mode 100644 index 000000000..1e5dbd757 --- /dev/null +++ b/server/scripts/benchmark_kimi_k3_deployments.py @@ -0,0 +1,573 @@ +#!/usr/bin/env python3 +"""Reproducible Kimi K3 SSD deployment comparison. + +The harness starts one server at a time, runs deterministic cold/warm requests, +captures the complete server log (including MoE NVMe shutdown telemetry), and +writes machine-readable JSON. It deliberately disables HTTP/prefix caches so +the only warm state under test is the routed-expert device cache. +""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import re +import signal +import socket +import subprocess +import sys +import time +import urllib.error +import urllib.request +from dataclasses import asdict, dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + + +_SPLIT_GGUF = re.compile(r"^(?P.+)-(?P\d+)-of-(?P\d+)(?P\.gguf)$") +_NVME_TELEMETRY = re.compile( + r"\[moe-nvme\] io=(?P\S+) requests=(?P\d+) reads=(?P\d+) " + r"payload=(?P[\d.]+) GiB physical=(?P[\d.]+) GiB " + r"active-io-rate=(?P[\d.]+) GiB/s cache-hit=(?P[\d.]+)% " + r"mean-demand-wait=(?P[\d.]+) ms .*?timeouts=(?P\d+) " + r"errors=(?P\d+) device-cache=(?P[\d.]+) MiB " + r"slots=(?P\d+) hits=(?P\d+) misses=(?P\d+) " + r"evictions=(?P\d+) graphs=(?P\d+) " + r"graph-hits=(?P\d+) graph-evictions=(?P\d+) " + r"launches=(?P\d+)" +) + + +@dataclass(frozen=True) +class Deployment: + name: str + primary_device: int + secondary_device: int | None + + +def build_deployments( + profiles: list[str], + strix_device: int, + r9700_device: int, + hetero_primary: str, +) -> list[Deployment]: + if strix_device < 0 or r9700_device < 0: + raise ValueError("GPU device indices must be non-negative") + if strix_device == r9700_device and "heterogeneous" in profiles: + raise ValueError("heterogeneous mode requires distinct R9700 and Strix devices") + + deployments: list[Deployment] = [] + for profile in profiles: + if profile == "strix-only": + deployments.append(Deployment("strix-only-ssd", strix_device, None)) + elif profile == "heterogeneous": + if hetero_primary == "strix": + deployments.append( + Deployment("heterogeneous-ssd", strix_device, r9700_device) + ) + else: + deployments.append( + Deployment("heterogeneous-ssd", r9700_device, strix_device) + ) + else: + raise ValueError(f"unknown deployment profile: {profile}") + return deployments + + +def deployment_environment( + base: dict[str, str], + deployment: Deployment, + nvme_backend: str, + primary_share_per_mille: int, + placement: Path | None, + device_cache_mb: int | None, + dual_trace: bool, +) -> dict[str, str]: + """Return a clean MoE environment without inherited benchmark tuning.""" + env = dict(base) + for key in list(env): + if key.startswith("DFLASH_MOE_NVME_") or key in { + "DFLASH_MOE_STORAGE", + "DFLASH_MOE_TP_GPU", + "DFLASH_MOE_PLACEMENT", + "DFLASH_MOE_PRIMARY_SHARE_PER_MILLE", + "DFLASH_MOE_DUAL_STREAM_TRACE", + }: + env.pop(key) + + env["DFLASH_MOE_NVME_BACKEND"] = nvme_backend + if device_cache_mb is not None: + env["DFLASH_MOE_NVME_DEVICE_CACHE_MB"] = str(device_cache_mb) + if deployment.secondary_device is not None: + env["DFLASH_MOE_TP_GPU"] = str(deployment.secondary_device) + env["DFLASH_MOE_PRIMARY_SHARE_PER_MILLE"] = str(primary_share_per_mille) + if placement is not None: + env["DFLASH_MOE_PLACEMENT"] = str(placement) + if dual_trace: + env["DFLASH_MOE_DUAL_STREAM_TRACE"] = "1" + return env + + +def server_command( + server_bin: Path, + model: Path, + deployment: Deployment, + port: int, + max_ctx: int, + extra_server_args: list[str], +) -> list[str]: + return [ + str(server_bin), + str(model), + "--host", + "127.0.0.1", + "--port", + str(port), + "--target-device", + f"hip:{deployment.primary_device}", + "--max-ctx", + str(max_ctx), + "--moe-storage", + "ssd", + "--prefix-cache-slots", + "0", + "--prefill-cache-slots", + "0", + "--disk-prefix-cache", + "off", + *extra_server_args, + ] + + +def discover_model_files(first_shard: Path) -> list[Path]: + """Validate and return a complete split GGUF in shard order.""" + match = _SPLIT_GGUF.match(first_shard.name) + if match is None: + if not first_shard.is_file(): + raise FileNotFoundError(first_shard) + return [first_shard] + + if int(match.group("index")) != 1: + raise ValueError(f"expected the first split GGUF shard, got {first_shard.name}") + + width = len(match.group("index")) + total_text = match.group("total") + total = int(total_text) + expected = [ + first_shard.with_name( + f"{match.group('prefix')}-{index:0{width}d}-of-{total_text}{match.group('suffix')}" + ) + for index in range(1, total + 1) + ] + missing = [path for path in expected if not path.is_file()] + if missing: + sample = ", ".join(path.name for path in missing[:3]) + raise FileNotFoundError( + f"split GGUF is incomplete: missing {len(missing)}/{total} shard(s): {sample}" + ) + return expected + + +def extract_nvme_telemetry(log_path: Path) -> list[dict[str, Any]]: + telemetry: list[dict[str, Any]] = [] + if not log_path.exists(): + return telemetry + for match in _NVME_TELEMETRY.finditer(log_path.read_text(errors="replace")): + row: dict[str, Any] = {"io": match.group("io")} + for key, value in match.groupdict().items(): + if key == "io": + continue + row[key] = float(value) if "." in value else int(value) + telemetry.append(row) + return telemetry + + +def ensure_port_available(port: int) -> None: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind(("127.0.0.1", port)) + except OSError as exc: + raise RuntimeError(f"benchmark port {port} is already in use") from exc + + +def wait_until_ready( + process: subprocess.Popen[bytes], + base_url: str, + timeout_s: float, +) -> float: + started = time.perf_counter() + deadline = started + timeout_s + next_report = started + 30.0 + while time.perf_counter() < deadline: + return_code = process.poll() + if return_code is not None: + raise RuntimeError(f"server exited during startup with status {return_code}") + try: + with urllib.request.urlopen(f"{base_url}/v1/models", timeout=2) as response: + if 200 <= response.status < 300: + return time.perf_counter() - started + except (urllib.error.URLError, ConnectionError, TimeoutError): + pass + now = time.perf_counter() + if now >= next_report: + print(f" still loading model ({now - started:.0f}s)", flush=True) + next_report = now + 30.0 + time.sleep(1.0) + raise TimeoutError(f"server did not become ready within {timeout_s:.0f}s") + + +def stream_chat( + base_url: str, + prompt: str, + max_tokens: int, + timeout_s: float, +) -> dict[str, Any]: + payload = { + "model": "dflash", + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, + "temperature": 0.0, + "stream": True, + } + request = urllib.request.Request( + f"{base_url}/v1/chat/completions", + data=json.dumps(payload).encode(), + headers={"Accept": "text/event-stream", "Content-Type": "application/json"}, + ) + started = time.perf_counter() + first_token_at: float | None = None + last_token_at: float | None = None + chunks = 0 + content: list[str] = [] + reasoning: list[str] = [] + usage: dict[str, Any] = {} + with urllib.request.urlopen(request, timeout=timeout_s) as response: + for raw_line in response: + line = raw_line.decode(errors="replace").strip() + if not line.startswith("data:"): + continue + encoded = line[5:].strip() + if encoded == "[DONE]": + break + try: + event = json.loads(encoded) + except json.JSONDecodeError: + continue + if event.get("usage"): + usage = event["usage"] + choices = event.get("choices") or [] + if not choices: + continue + delta = choices[0].get("delta") or {} + visible = delta.get("content") or "" + thought = delta.get("reasoning_content") or "" + if not visible and not thought: + continue + now = time.perf_counter() + if first_token_at is None: + first_token_at = now + last_token_at = now + chunks += 1 + content.append(visible) + reasoning.append(thought) + + finished = time.perf_counter() + wall_s = finished - started + client_ttft_s = (first_token_at - started) if first_token_at is not None else wall_s + completion_tokens = int(usage.get("completion_tokens") or chunks) + event_decode_s = ( + last_token_at - first_token_at + if first_token_at is not None and last_token_at is not None + else 0.0 + ) + timings = usage.get("timings") or {} + server_prefill_s = float(timings.get("prefill_ms") or 0.0) / 1000.0 + server_decode_s = float(timings.get("decode_ms") or 0.0) / 1000.0 + server_decode_tok_s = float(timings.get("decode_tokens_per_sec") or 0.0) + if server_decode_tok_s <= 0.0: + server_decode_tok_s = ( + (completion_tokens - 1) / event_decode_s + if completion_tokens > 1 and event_decode_s > 0 + else 0.0 + ) + wall_tok_s = completion_tokens / wall_s if wall_s > 0 else 0.0 + event_decode_tok_s = ( + (completion_tokens - 1) / event_decode_s + if completion_tokens > 1 and event_decode_s > 0 + else 0.0 + ) + return { + "wall_s": wall_s, + "wall_tok_s": wall_tok_s, + "client_ttft_s": client_ttft_s, + "server_prefill_s": server_prefill_s, + "server_decode_s": server_decode_s, + "server_decode_tok_s": server_decode_tok_s, + "event_decode_s": event_decode_s, + "event_decode_tok_s": event_decode_tok_s, + "completion_tokens": completion_tokens, + "event_chunks": chunks, + "content": "".join(content), + "reasoning_content": "".join(reasoning), + "usage": usage, + } + + +def stop_server(process: subprocess.Popen[bytes], timeout_s: float = 60.0) -> int: + if process.poll() is None: + process.send_signal(signal.SIGINT) + try: + return process.wait(timeout=timeout_s) + except subprocess.TimeoutExpired: + process.terminate() + try: + return process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + return process.wait() + + +def selected_environment(env: dict[str, str]) -> dict[str, str]: + return { + key: value + for key, value in sorted(env.items()) + if key.startswith("DFLASH_MOE_") + } + + +def run_deployment( + args: argparse.Namespace, + deployment: Deployment, + output_dir: Path, +) -> dict[str, Any]: + ensure_port_available(args.port) + env = deployment_environment( + os.environ, + deployment, + args.nvme_backend, + args.primary_share_per_mille, + args.placement, + args.device_cache_mb, + args.dual_trace, + ) + command = server_command( + args.server_bin, + args.model, + deployment, + args.port, + args.max_ctx, + args.extra_server_arg, + ) + log_path = output_dir / f"{deployment.name}.server.log" + result: dict[str, Any] = { + "deployment": asdict(deployment), + "command": command, + "environment": selected_environment(env), + "server_log": str(log_path), + "requests": [], + } + print( + f"\n[{deployment.name}] primary=hip:{deployment.primary_device} " + f"secondary={deployment.secondary_device}", + flush=True, + ) + process: subprocess.Popen[bytes] | None = None + with log_path.open("wb") as log_file: + try: + process = subprocess.Popen( + command, + env=env, + stdout=log_file, + stderr=subprocess.STDOUT, + ) + result["startup_s"] = wait_until_ready( + process, f"http://127.0.0.1:{args.port}", args.startup_timeout + ) + print(f" server ready in {result['startup_s']:.1f}s", flush=True) + for iteration in range(args.iterations): + measurement = stream_chat( + f"http://127.0.0.1:{args.port}", + args.prompt, + args.max_tokens, + args.request_timeout, + ) + measurement["state"] = "cold" if iteration == 0 else "warm" + result["requests"].append(measurement) + print( + f" {measurement['state']}[{iteration}] " + f"wall={measurement['wall_s']:.3f}s " + f"prefill={measurement['server_prefill_s']:.3f}s " + f"decode={measurement['server_decode_tok_s']:.3f} tok/s", + flush=True, + ) + except Exception as exc: # Preserve the other profile and its evidence. + result["error"] = f"{type(exc).__name__}: {exc}" + print(f" FAILED: {result['error']}", file=sys.stderr, flush=True) + finally: + if process is not None: + result["server_exit_code"] = stop_server(process) + result["nvme_telemetry"] = extract_nvme_telemetry(log_path) + outputs = [ + row["reasoning_content"] + row["content"] for row in result["requests"] + ] + result["within_profile_output_match"] = bool(outputs) and len(set(outputs)) == 1 + (output_dir / f"{deployment.name}.json").write_text( + json.dumps(result, indent=2, sort_keys=True) + "\n" + ) + return result + + +def create_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Compare Kimi K3 Strix-only and heterogeneous SSD deployments" + ) + parser.add_argument("model", type=Path, help="first shard of the split Kimi K3 GGUF") + parser.add_argument( + "--server-bin", type=Path, default=Path("server/build-hip-dual/dflash_server") + ) + parser.add_argument( + "--profiles", + nargs="+", + choices=("strix-only", "heterogeneous"), + default=("strix-only", "heterogeneous"), + ) + parser.add_argument("--strix-device", type=int, default=1) + parser.add_argument("--r9700-device", type=int, default=0) + parser.add_argument( + "--hetero-primary", + choices=("strix", "r9700"), + default="strix", + help=( + "primary device for the heterogeneous run; current IQ1_S defaults to Strix " + "because its non-routed tensors exceed R9700 VRAM" + ), + ) + parser.add_argument("--primary-share-per-mille", type=int, default=500) + parser.add_argument("--placement", type=Path) + parser.add_argument("--device-cache-mb", type=int) + parser.add_argument( + "--nvme-backend", choices=("auto", "uring", "pread", "mmap"), default="auto" + ) + parser.add_argument("--port", type=int, default=18080) + parser.add_argument("--max-ctx", type=int, default=8192) + parser.add_argument("--max-tokens", type=int, default=8) + parser.add_argument("--iterations", type=int, default=3) + parser.add_argument("--startup-timeout", type=float, default=3600.0) + parser.add_argument("--request-timeout", type=float, default=3600.0) + parser.add_argument( + "--prompt", + default="Explain in two short sentences why the sky appears blue.", + ) + parser.add_argument("--dual-trace", action="store_true") + parser.add_argument("--extra-server-arg", action="append", default=[]) + parser.add_argument("--output-dir", type=Path) + parser.add_argument( + "--dry-run", action="store_true", help="print resolved commands without starting servers" + ) + return parser + + +def validate_args(args: argparse.Namespace) -> None: + if not 1 <= args.port <= 65535: + raise ValueError("port must be in [1, 65535]") + if args.max_ctx <= 0 or args.max_tokens <= 0 or args.iterations <= 0: + raise ValueError("max-ctx, max-tokens, and iterations must be positive") + if not 0 <= args.primary_share_per_mille <= 1000: + raise ValueError("primary-share-per-mille must be in [0, 1000]") + if args.device_cache_mb is not None and args.device_cache_mb < 0: + raise ValueError("device-cache-mb must be non-negative") + if len(set(args.profiles)) != len(args.profiles): + raise ValueError("deployment profiles must not be repeated") + if args.placement is not None and not args.dry_run and not args.placement.is_file(): + raise FileNotFoundError(args.placement) + + +def main() -> int: + parser = create_parser() + args = parser.parse_args() + try: + validate_args(args) + deployments = build_deployments( + list(args.profiles), + args.strix_device, + args.r9700_device, + args.hetero_primary, + ) + if args.dry_run: + resolved = [] + for deployment in deployments: + env = deployment_environment( + os.environ, + deployment, + args.nvme_backend, + args.primary_share_per_mille, + args.placement, + args.device_cache_mb, + args.dual_trace, + ) + resolved.append( + { + "deployment": asdict(deployment), + "command": server_command( + args.server_bin, + args.model, + deployment, + args.port, + args.max_ctx, + args.extra_server_arg, + ), + "environment": selected_environment(env), + } + ) + print(json.dumps(resolved, indent=2, sort_keys=True)) + return 0 + + if not args.server_bin.is_file(): + raise FileNotFoundError(args.server_bin) + model_files = discover_model_files(args.model) + except (FileNotFoundError, RuntimeError, ValueError) as exc: + parser.error(str(exc)) + + timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") + output_dir = args.output_dir or Path("bench-out") / f"kimi-k3-{timestamp}" + output_dir.mkdir(parents=True, exist_ok=False) + manifest = { + "created_at": datetime.now(UTC).isoformat(), + "hostname": platform.node(), + "platform": platform.platform(), + "model_first_shard": str(args.model), + "model_shards": len(model_files), + "model_bytes": sum(path.stat().st_size for path in model_files), + "results": [], + } + print(f"Writing benchmark evidence to {output_dir}", flush=True) + for deployment in deployments: + manifest["results"].append(run_deployment(args, deployment, output_dir)) + + successful = [row for row in manifest["results"] if row.get("requests")] + outputs = [ + row["requests"][0]["reasoning_content"] + row["requests"][0]["content"] + for row in successful + ] + manifest["cross_profile_output_match"] = bool(outputs) and len(set(outputs)) == 1 + manifest_path = output_dir / "comparison.json" + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") + print(f"\nComparison: {manifest_path}") + + if any("error" in row for row in manifest["results"]): + return 1 + deterministic = all(row["within_profile_output_match"] for row in successful) + if len(successful) > 1: + deterministic = deterministic and manifest["cross_profile_output_match"] + if not deterministic: + print("ERROR: deterministic outputs differ within or across profiles", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/server/scripts/test_benchmark_kimi_k3_deployments.py b/server/scripts/test_benchmark_kimi_k3_deployments.py new file mode 100644 index 000000000..6162115be --- /dev/null +++ b/server/scripts/test_benchmark_kimi_k3_deployments.py @@ -0,0 +1,109 @@ +import tempfile +import unittest +from pathlib import Path + +from benchmark_kimi_k3_deployments import ( + Deployment, + build_deployments, + deployment_environment, + discover_model_files, + extract_nvme_telemetry, + server_command, +) + + +class KimiDeploymentBenchmarkTests(unittest.TestCase): + def test_builds_capacity_safe_default_profiles(self): + profiles = build_deployments(["strix-only", "heterogeneous"], 1, 0, "strix") + self.assertEqual( + profiles, + [ + Deployment("strix-only-ssd", 1, None), + Deployment("heterogeneous-ssd", 1, 0), + ], + ) + + def test_r9700_primary_is_explicit(self): + profiles = build_deployments(["heterogeneous"], 1, 0, "r9700") + self.assertEqual(profiles, [Deployment("heterogeneous-ssd", 0, 1)]) + with self.assertRaisesRegex(ValueError, "distinct"): + build_deployments(["heterogeneous"], 0, 0, "strix") + + def test_environment_removes_stale_tuning(self): + base = { + "PATH": "/bin", + "DFLASH_MOE_STORAGE": "resident", + "DFLASH_MOE_NVME_SLOTS": "64", + "DFLASH_MOE_TP_GPU": "9", + "DFLASH_MOE_PLACEMENT": "/stale.json", + } + env = deployment_environment( + base, + Deployment("heterogeneous-ssd", 1, 0), + "uring", + 600, + Path("/new.json"), + 4096, + True, + ) + self.assertEqual(env["PATH"], "/bin") + self.assertNotIn("DFLASH_MOE_STORAGE", env) + self.assertNotIn("DFLASH_MOE_NVME_SLOTS", env) + self.assertEqual(env["DFLASH_MOE_NVME_BACKEND"], "uring") + self.assertEqual(env["DFLASH_MOE_NVME_DEVICE_CACHE_MB"], "4096") + self.assertEqual(env["DFLASH_MOE_TP_GPU"], "0") + self.assertEqual(env["DFLASH_MOE_PRIMARY_SHARE_PER_MILLE"], "600") + self.assertEqual(env["DFLASH_MOE_PLACEMENT"], "/new.json") + self.assertEqual(env["DFLASH_MOE_DUAL_STREAM_TRACE"], "1") + + def test_strix_only_has_no_secondary_owner(self): + env = deployment_environment( + {"DFLASH_MOE_TP_GPU": "0"}, + Deployment("strix-only-ssd", 1, None), + "auto", + 500, + None, + None, + False, + ) + self.assertNotIn("DFLASH_MOE_TP_GPU", env) + command = server_command( + Path("server"), Path("model.gguf"), Deployment("strix", 1, None), 8080, 8192, [] + ) + self.assertIn("hip:1", command) + self.assertEqual(command[command.index("--moe-storage") + 1], "ssd") + self.assertEqual(command[command.index("--prefix-cache-slots") + 1], "0") + + def test_discovers_complete_split_model(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + paths = [root / f"model-{index:05d}-of-00003.gguf" for index in range(1, 4)] + for index, path in enumerate(paths, 1): + path.write_bytes(bytes(index)) + self.assertEqual(discover_model_files(paths[0]), paths) + with self.assertRaisesRegex(ValueError, "first"): + discover_model_files(paths[1]) + paths[-1].unlink() + with self.assertRaisesRegex(FileNotFoundError, "incomplete"): + discover_model_files(paths[0]) + + def test_extracts_each_owner_telemetry_line(self): + line = ( + "[moe-nvme] io=io_uring requests=10 reads=9 payload=1.250 GiB " + "physical=1.500 GiB active-io-rate=3.750 GiB/s cache-hit=10.0% " + "mean-demand-wait=2.500 ms dedupe=0 upgrades=0 dropped-prefetch=0 " + "timeouts=0 errors=0 device-cache=1024.0 MiB slots=2 hits=1 misses=9 " + "evictions=3 graphs=1 graph-hits=8 graph-evictions=0 launches=9\n" + ) + with tempfile.TemporaryDirectory() as directory: + log = Path(directory) / "server.log" + log.write_text(line + line) + rows = extract_nvme_telemetry(log) + self.assertEqual(len(rows), 2) + self.assertEqual(rows[0]["io"], "io_uring") + self.assertEqual(rows[0]["active_io_gib_s"], 3.75) + self.assertEqual(rows[0]["errors"], 0) + + +if __name__ == "__main__": + unittest.main() From e0a0b4e6789646658ee6e8be915ac30174087915 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sat, 1 Aug 2026 02:08:26 +0200 Subject: [PATCH 14/20] feat(moe): add expert-major package and profile cache --- .github/workflows/ci.yml | 15 +- server/CMakeLists.txt | 16 + server/docs/ENVIRONMENT.md | 6 + server/docs/MOE_NVME_STREAMING.md | 62 +- server/src/common/moe_expert_package.cpp | 730 ++++++++++++++++++ server/src/common/moe_expert_package.h | 84 ++ server/src/common/moe_hybrid_stream.cpp | 146 +++- server/src/common/moe_hybrid_stream.h | 31 + server/src/common/moe_stream_cache_policy.cpp | 91 +++ server/src/common/moe_stream_cache_policy.h | 42 + server/src/kimi_k3/kimi_k3_backend.cpp | 540 +++++++++++-- server/src/kimi_k3/kimi_k3_backend.h | 5 + server/src/kimi_k3/kimi_k3_graph.cpp | 17 +- server/src/kimi_k3/kimi_k3_internal.h | 5 +- server/test/test_moe_expert_package.cpp | 263 +++++++ server/test/test_moe_stream_compute.cpp | 71 ++ .../test/test_moe_stream_owner_partition.cpp | 51 ++ 17 files changed, 2064 insertions(+), 111 deletions(-) create mode 100644 server/src/common/moe_expert_package.cpp create mode 100644 server/src/common/moe_expert_package.h create mode 100644 server/src/common/moe_stream_cache_policy.cpp create mode 100644 server/src/common/moe_stream_cache_policy.h create mode 100644 server/test/test_moe_expert_package.cpp diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 00804a9ab..28a0ca994 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -79,13 +79,14 @@ jobs: -DCMAKE_BUILD_TYPE=Release cmake --build build --target \ test_dflash test_generate test_flash_attn_sparse test_server_unit \ - test_deepseek4_unit test_moe_nvme_scheduler -j$(nproc) + test_deepseek4_unit test_moe_nvme_scheduler \ + test_moe_expert_package -j$(nproc) - name: Run C++ server unit tests run: | cd server/build ctest --output-on-failure \ - -R "server_unit|deepseek4_unit|test_moe_nvme_scheduler" \ + -R "server_unit|deepseek4_unit|test_moe_nvme_scheduler|test_moe_expert_package" \ --no-tests=error - name: Populate venv with cu128 torch + setuptools @@ -177,6 +178,7 @@ jobs: cmake --build build \ --target test_flash_attn_sparse test_deepseek4_mmid_grouped_cuda \ test_moe_stream_compute test_moe_nvme_scheduler \ + test_moe_expert_package \ -j"$(nproc)" - name: Run flash-attn sparse kernel test on the 3090 @@ -191,7 +193,9 @@ jobs: run: ./server/build/test_moe_stream_compute - name: Run NVMe scheduler production tests - run: ./server/build/test_moe_nvme_scheduler + run: | + ./server/build/test_moe_nvme_scheduler + ./server/build/test_moe_expert_package # Optional model-backed end-to-end smoke (real spec-decode on the 3090), # disabled by default because it builds dflash_server and lazy-loads the @@ -300,11 +304,12 @@ jobs: cmake --build "$RUNNER_TEMP/rocmfp-build" \ --target test_rocmfp4 test_rocmfpx test_rocmfp4_hip_tail test_rocmfpx_mmq \ test_deepseek4_mmid_grouped_cuda test_moe_stream_compute \ - test_moe_nvme_scheduler test_recurrent_snapshot test_server_unit \ + test_moe_nvme_scheduler test_moe_expert_package \ + test_recurrent_snapshot test_server_unit \ --parallel 8 ctest --test-dir "$RUNNER_TEMP/rocmfp-build" \ --output-on-failure \ - -R 'rocmfp4_reference|rocmfpx_reference|rocmfp4_hip_tail|rocmfpx_mmq|deepseek4_mmid_grouped_cuda|test_moe_stream_compute|test_moe_nvme_scheduler|recurrent_snapshot|ChainRollbackPolicy' + -R 'rocmfp4_reference|rocmfpx_reference|rocmfp4_hip_tail|rocmfpx_mmq|deepseek4_mmid_grouped_cuda|test_moe_stream_compute|test_moe_nvme_scheduler|test_moe_expert_package|recurrent_snapshot|ChainRollbackPolicy' build-windows: name: Build Windows (MSVC + CUDA, library + server targets) diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 745b9717e..c3acbbe0f 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -312,6 +312,8 @@ add_library(dflash_common STATIC src/common/spark_corpus.cpp src/common/moe_hybrid_ffn_eval.cpp src/common/moe_nvme_scheduler.cpp + src/common/moe_expert_package.cpp + src/common/moe_stream_cache_policy.cpp src/common/moe_hybrid_stream.cpp src/common/moe_expert_compute.cpp src/common/moe_expert_compute_cpu.cpp @@ -695,6 +697,20 @@ if(DFLASH27B_TESTS) list(APPEND _raw_unit_test_targets test_moe_nvme_scheduler) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_moe_expert_package.cpp") + add_executable(test_moe_expert_package + test/test_unit_main.cpp + test/test_moe_expert_package.cpp + src/common/moe_expert_package.cpp + src/common/moe_nvme_scheduler.cpp) + target_include_directories(test_moe_expert_package PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/src/common + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + target_link_libraries(test_moe_expert_package PRIVATE Threads::Threads) + list(APPEND _raw_unit_test_targets test_moe_expert_package) + endif() + # Read-only microbenchmark for tuning the exact expert I/O path on the # deployment SSD. It accepts any sufficiently large file; no model parser # or generated benchmark data is required. diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index 5a62c9145..0e92daf23 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -35,6 +35,8 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_MOE_STORAGE` | auto | Routed-MoE storage policy (`auto`, `resident`, `ssd`); prefer `--moe-storage`, which takes precedence. | | `DFLASH_MOE_NVME_COLD_TIER` | unset | DEPRECATED: DeepSeek compatibility alias (`auto`, `on`, `off`). | | `DFLASH_MOE_NVME_*` | tuned defaults | BURN-IN: bounded MoE SSD scheduler/backend controls; see `MOE_NVME_STREAMING.md`. | +| `DFLASH_MOE_EXPERT_PACKAGE` / `DFLASH_MOE_EXPERT_PACKAGE_BUILD` | unset | BURN-IN: use or explicitly build a validated, byte-exact expert-major SSD package. | +| `DFLASH_MOE_ROUTE_STATS_OUT` / `DFLASH_MOE_HOTNESS_CSV` | unset | BURN-IN: capture native route counts, then profile-warm and pin valuable streamed experts. | | `GGML_CUDA_BATCH_PEER_COPIES` | unset | BURN-IN: publish ordered HIP peer copies with one cross-device dependency per source/destination pair. | | `DFLASH_MOE_PREFILL_PERSISTENT_OWNER_ALLOC` | 1 for qualified long heterogeneous prefill | KILL SWITCH: =0 restores per-layer route/owner scratch allocation. | | `DFLASH_MOE_TP_*` / `DFLASH_MOE_HYBRID_PREFILL_EAGER` | unset | BURN-IN: model-neutral names for common heterogeneous-MoE scheduling and kernel policy. Existing `DFLASH_DS4_*` names remain compatibility aliases. | @@ -175,18 +177,22 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH_MOE_EXPERT_COMPUTE_IPC_SHARED_BYTES` - moe_expert_compute_ipc.cpp - `DFLASH_MOE_EXPERT_COMPUTE_IPC_TRANSPORT` - moe_expert_compute_ipc.cpp - `DFLASH_MOE_EXPERT_COMPUTE_THREADS` - moe_expert_compute_cpu.cpp +- `DFLASH_MOE_EXPERT_PACKAGE` - kimi_k3_backend.cpp +- `DFLASH_MOE_EXPERT_PACKAGE_BUILD` - kimi_k3_backend.cpp - `DFLASH_MOE_EXPERT_MAJOR_GPU_REDUCE` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_EXPERT_MAJOR_PREFILL` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_FIXED_SLOT_GRAPHS` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_FIXED_SLOT_MAX` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_FULL_COLD_PARALLEL` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_FUSED_COMBINE` - moe_hybrid_ffn_eval.cpp +- `DFLASH_MOE_HOTNESS_CSV` - kimi_k3_backend.cpp - `DFLASH_MOE_PLACEMENT` - kimi_k3_backend.cpp - `DFLASH_MOE_PREFILL_DEVICE_INPUT` - deepseek4_graph.cpp - `DFLASH_MOE_PREFILL_HOT_SUB_BATCH` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_PREFILL_MASKED_COLD` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_PREFILL_PERSISTENT_OWNER_ALLOC` - deepseek4_graph.cpp - `DFLASH_MOE_PRIMARY_SHARE_PER_MILLE` - moe_hybrid_stream.cpp +- `DFLASH_MOE_ROUTE_STATS_OUT` - kimi_k3_backend.cpp - `DFLASH_NO_MASK` - laguna_backend.cpp - `DFLASH_NO_MOE_ROUTER_FUSE` - qwen35moe_ffn.cpp - `DFLASH_NO_MOE_SWIGLU_FUSE` - qwen35moe_ffn.cpp diff --git a/server/docs/MOE_NVME_STREAMING.md b/server/docs/MOE_NVME_STREAMING.md index 229c0b01b..4d732597e 100644 --- a/server/docs/MOE_NVME_STREAMING.md +++ b/server/docs/MOE_NVME_STREAMING.md @@ -173,6 +173,50 @@ Kimi's native router remains authoritative. The current text backend is correctness-first and sequential; captured per-layer graphs and the vision tower are separate optimizations. +### Expert-major package and profile-guided cache + +Standard split GGUF stores gate/up/down expert tensors in separate ranges, so +one cache miss normally needs three SSD reads. Kimi can compile the routed +bytes once into a self-describing expert-major package. This is a byte-exact +layout transform: it changes neither quantization nor model output, and one +aligned expert record replaces the three reads. + +```bash +export DFLASH_MOE_EXPERT_PACKAGE=/fast-nvme/Kimi-K3.lbmoe +export DFLASH_MOE_EXPERT_PACKAGE_BUILD=1 + +./build-hip/dflash_server \ + /models/Kimi-K3-UD-IQ1_S-00001-of-00014.gguf \ + --target-device hip:0 --max-ctx 8192 +``` + +The one-time build streams source shards and reports completed layers. It +needs free disk roughly equal to the routed expert weights. A process-unique +temporary file is synced and atomically published only after its package magic +is valid, so an interrupted build never replaces a working artifact. On later +starts, unset `DFLASH_MOE_EXPERT_PACKAGE_BUILD`; `force` explicitly rebuilds +an existing artifact. The loader validates model shape and source-layout +fingerprint plus deterministic weight samples before using it, so a same-shape +package from another checkpoint also fails closed. + +For stable workloads, a short representative run can record native-router +counts, and the next run can pin the highest observed value-per-byte experts: + +```bash +# Profiling run: ordinary adaptive cache, native routes unchanged. +export DFLASH_MOE_ROUTE_STATS_OUT=/models/kimi-routes.csv + +# Deployment run: warm and pin the profile-selected experts. +unset DFLASH_MOE_ROUTE_STATS_OUT +export DFLASH_MOE_HOTNESS_CSV=/models/kimi-routes.csv +``` + +At least one quarter of each device cache (and never fewer than two slots) +remains adaptive for profile drift and misses. With two GPUs, the same runtime +ownership rule filters both warm plans, so an expert is never pinned on both +R9700 and Strix. A missing profile simply retains the adaptive LFRU cache; +prediction is not required for correctness. + On a two-GPU Lucebox, put the compute-intensive primary path on the R9700 and use Strix as the secondary capacity owner. Both devices receive independent SSD slots/caches and execute their selected routed experts concurrently: @@ -216,6 +260,10 @@ not improve the qualified P310 drive and consumes extra pinned/system memory. | `DFLASH_MOE_NVME_DEVICE_CACHE_MB` | automatic/`0` | Override adaptive device expert-cache memory; `0` leaves only pipeline slots | | `DFLASH_MOE_NVME_GRAPH_CACHE` | `8` | Persistent expert-graph variants retained per stream engine; `0` is a diagnostic no-cache mode | | `DFLASH_MOE_NVME_REFERENCE_EVAL` | unset | Diagnostic only: `1` restores the allocation-heavy reference evaluator for numerical/performance A/B | +| `DFLASH_MOE_EXPERT_PACKAGE` | unset | Optional validated expert-major package used in place of tensor-major GGUF expert regions | +| `DFLASH_MOE_EXPERT_PACKAGE_BUILD` | unset | One-time Kimi package creation: `1` builds when absent; `force` rebuilds | +| `DFLASH_MOE_ROUTE_STATS_OUT` | unset | Record Kimi native-router counts as a reusable model-neutral CSV profile | +| `DFLASH_MOE_HOTNESS_CSV` | unset | Warm and pin the highest-value experts from a compatible routing profile | | `DFLASH_MOE_TP_GPU` | primary GPU | Optional secondary GPU; enables concurrent route ownership when different from the primary | | `DFLASH_MOE_PLACEMENT` | unset | Offline placement JSON; listed experts belong to the primary GPU | | `DFLASH_MOE_PRIMARY_SHARE_PER_MILLE` | `500` | Bring-up hash split used only when no placement is supplied | @@ -230,7 +278,8 @@ only when it covers the complete logical payload at an unaligned file tail. tiny experts and checks both tensor-major and expert-major GPU results against a CPU oracle. It defaults to GPU 0, so it runs directly on Strix-only systems; `DFLASH_TEST_GPU` selects another device on multi-GPU hosts. The standalone -targets `test_moe_nvme_scheduler`, `bench_moe_nvme_io`, and +targets `test_moe_nvme_scheduler`, `test_moe_expert_package`, +`bench_moe_nvme_io`, and `bench_moe_nvme_pipeline` test scheduling, raw storage, and the complete SSD-to-GPU path. Benchmarks are read-only. @@ -303,8 +352,9 @@ that leaves the native router untouched. MoE-SpAc motivates compile-time expert layout and I/O coalescing. Tutti's slack-aware `io_uring` scheduling is relevant when persistent KV traffic shares the device. -The persistent graph, model-neutral evaluator, and exact concurrent owner -split are now shared infrastructure. The next high-value work is a repacked -expert-major artifact and a device-resident fork/join that removes Kimi's host -activation boundary. Any learned predictor comes after that deterministic path -is qualified, and may only prefetch routes selected later by the native router. +The persistent graph, model-neutral evaluator, exact concurrent owner split, +expert-major package, and profile-guided pinned cache are now shared +infrastructure. The next high-value work is a device-resident fork/join that +removes Kimi's host activation boundary. Any learned predictor comes after +that deterministic path is qualified, and may only prefetch routes selected +later by the native router. diff --git a/server/src/common/moe_expert_package.cpp b/server/src/common/moe_expert_package.cpp new file mode 100644 index 000000000..09ec4c665 --- /dev/null +++ b/server/src/common/moe_expert_package.cpp @@ -0,0 +1,730 @@ +#include "moe_expert_package.h" + +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#else +#include +#endif + +namespace dflash::common { +namespace { + +constexpr uint8_t kMagic[8] = {'L', 'B', 'M', 'O', 'E', 'P', 'K', '1'}; +constexpr uint32_t kVersion = 1; +constexpr size_t kFixedHeaderBytes = 64; +constexpr size_t kLayerEntryBytes = 128; +constexpr size_t kCopyChunkBytes = 8 * 1024 * 1024; + +bool checked_add(size_t a, size_t b, size_t & out) { + if (a > std::numeric_limits::max() - b) return false; + out = a + b; + return true; +} + +bool checked_mul(size_t a, size_t b, size_t & out) { + if (a != 0 && b > std::numeric_limits::max() / a) return false; + out = a * b; + return true; +} + +bool is_power_of_two(size_t value) { + return value != 0 && (value & (value - 1)) == 0; +} + +bool align_up(size_t value, size_t alignment, size_t & out) { + if (!is_power_of_two(alignment)) return false; + const size_t mask = alignment - 1; + if (value > std::numeric_limits::max() - mask) return false; + out = (value + mask) & ~mask; + return true; +} + +void put_u32(uint8_t * dst, uint32_t value) { + for (int i = 0; i < 4; ++i) dst[i] = (uint8_t) (value >> (8 * i)); +} + +void put_u64(uint8_t * dst, uint64_t value) { + for (int i = 0; i < 8; ++i) dst[i] = (uint8_t) (value >> (8 * i)); +} + +uint32_t get_u32(const uint8_t * src) { + uint32_t value = 0; + for (int i = 0; i < 4; ++i) value |= (uint32_t) src[i] << (8 * i); + return value; +} + +uint64_t get_u64(const uint8_t * src) { + uint64_t value = 0; + for (int i = 0; i < 8; ++i) value |= (uint64_t) src[i] << (8 * i); + return value; +} + +void hash_u64(uint64_t & hash, uint64_t value) { + constexpr uint64_t kPrime = 1099511628211ULL; + for (int i = 0; i < 8; ++i) { + hash ^= (uint8_t) (value >> (8 * i)); + hash *= kPrime; + } +} + +void hash_bytes(uint64_t & hash, const uint8_t * data, size_t bytes) { + constexpr uint64_t kPrime = 1099511628211ULL; + for (size_t i = 0; i < bytes; ++i) { + hash ^= data[i]; + hash *= kPrime; + } +} + +bool read_at(const MoeNvmeSource & source, size_t offset, + void * dst, size_t bytes, std::string * err) { + if (offset > source.mmap_size || bytes > source.mmap_size - offset) { + if (err) *err = "expert package source read is out of bounds"; + return false; + } + if (bytes == 0) return true; + if (source.mmap_data) { + std::memcpy(dst, + static_cast(source.mmap_data) + offset, + bytes); + return true; + } + if (source.fd < 0) { + if (err) *err = "expert package source has neither mmap data nor fd"; + return false; + } + + size_t done = 0; + while (done < bytes) { + const size_t chunk = std::min(bytes - done, kCopyChunkBytes); +#if defined(_WIN32) + if (::_lseeki64(source.fd, (__int64) (offset + done), SEEK_SET) < 0) { + if (err) *err = "expert package source seek failed"; + return false; + } + const int result = ::_read(source.fd, + static_cast(dst) + done, + (unsigned int) chunk); +#else + const ssize_t result = ::pread(source.fd, + static_cast(dst) + done, + chunk, (off_t) (offset + done)); +#endif + if (result < 0 && errno == EINTR) continue; + if (result <= 0) { + if (err) *err = std::string("expert package source read failed: ") + + std::strerror(errno); + return false; + } + done += (size_t) result; + } + return true; +} + +bool write_at(int fd, size_t offset, const void * src, + size_t bytes, std::string * err) { + size_t done = 0; + while (done < bytes) { + const size_t chunk = std::min(bytes - done, kCopyChunkBytes); +#if defined(_WIN32) + if (::_lseeki64(fd, (__int64) (offset + done), SEEK_SET) < 0) { + if (err) *err = "expert package output seek failed"; + return false; + } + const int result = ::_write(fd, + static_cast(src) + done, + (unsigned int) chunk); +#else + const ssize_t result = ::pwrite(fd, + static_cast(src) + done, + chunk, (off_t) (offset + done)); +#endif + if (result < 0 && errno == EINTR) continue; + if (result <= 0) { + if (err) *err = std::string("expert package output write failed: ") + + std::strerror(errno); + return false; + } + done += (size_t) result; + } + return true; +} + +bool resize_file(int fd, size_t bytes, std::string * err) { +#if defined(_WIN32) + if (::_chsize_s(fd, bytes) == 0) return true; +#else + if (::ftruncate(fd, (off_t) bytes) == 0) return true; +#endif + if (err) *err = std::string("failed to size expert package: ") + + std::strerror(errno); + return false; +} + +bool sync_file(int fd, std::string * err) { +#if defined(_WIN32) + if (::_commit(fd) == 0) return true; +#else + if (::fsync(fd) == 0) return true; +#endif + if (err) *err = std::string("failed to sync expert package: ") + + std::strerror(errno); + return false; +} + +bool validate_component(const ExpertFileRegion & region, + size_t expert_bytes, uint32_t experts, + const std::vector & sources, + const char * label, std::string * err) { + if (expert_bytes == 0 || experts == 0 || + region.source_index >= sources.size()) { + if (err) *err = std::string("invalid expert package ") + label + + " component"; + return false; + } + size_t required = 0; + if (!checked_mul(expert_bytes, (size_t) experts, required) || + region.size < required || + region.offset > sources[region.source_index].mmap_size || + required > sources[region.source_index].mmap_size - region.offset) { + if (err) *err = std::string("expert package ") + label + + " component exceeds its source"; + return false; + } + return true; +} + +struct ComponentRef { + const ExpertFileRegion * source = nullptr; + size_t bytes = 0; + size_t destination_offset = 0; +}; + +int layer_components(const LayerExpertRegions & layer, + ComponentRef (&components)[3]) { + if (layer.fused_gate_up) { + components[0] = {&layer.gate_up_exps, + layer.expert_bytes_gate_up, + layer.expert_major.gate_up_offset}; + components[1] = {&layer.down_exps, + layer.expert_bytes_down, + layer.expert_major.down_offset}; + return 2; + } + components[0] = {&layer.gate_exps, + layer.expert_bytes_gate, + layer.expert_major.gate_offset}; + components[1] = {&layer.up_exps, + layer.expert_bytes_up, + layer.expert_major.up_offset}; + components[2] = {&layer.down_exps, + layer.expert_bytes_down, + layer.expert_major.down_offset}; + return 3; +} + +bool encode_header(const MoeExpertPackageManifest & manifest, + std::vector & header, std::string * err) { + size_t table_bytes = 0; + if (!checked_mul(manifest.layer_regions.size(), kLayerEntryBytes, + table_bytes)) { + if (err) *err = "expert package header size overflow"; + return false; + } + size_t raw_header_bytes = 0; + if (!checked_add(kFixedHeaderBytes, table_bytes, raw_header_bytes)) { + if (err) *err = "expert package header size overflow"; + return false; + } + size_t header_bytes = 0; + if (!align_up(raw_header_bytes, manifest.record_alignment, header_bytes)) { + if (err) *err = "expert package header alignment overflow"; + return false; + } + if (header_bytes > std::numeric_limits::max()) { + if (err) *err = "expert package header exceeds its on-disk field"; + return false; + } + try { + header.assign(header_bytes, 0); + } catch (const std::bad_alloc &) { + if (err) *err = "failed to allocate expert package header"; + return false; + } + std::memcpy(header.data(), kMagic, sizeof(kMagic)); + put_u32(header.data() + 8, manifest.version); + put_u32(header.data() + 12, (uint32_t) header_bytes); + put_u32(header.data() + 16, (uint32_t) manifest.layer_regions.size()); + put_u32(header.data() + 20, (uint32_t) kLayerEntryBytes); + put_u64(header.data() + 24, manifest.file_bytes); + put_u64(header.data() + 32, manifest.source_layout_hash); + put_u64(header.data() + 40, manifest.record_alignment); + put_u64(header.data() + 48, manifest.component_alignment); + + for (size_t i = 0; i < manifest.layer_regions.size(); ++i) { + const LayerExpertRegions & layer = manifest.layer_regions[i]; + uint8_t * entry = header.data() + kFixedHeaderBytes + + i * kLayerEntryBytes; + put_u32(entry + 0, manifest.expert_counts[i]); + put_u32(entry + 4, layer.fused_gate_up ? 1U : 0U); + put_u64(entry + 8, layer.expert_major.experts.offset); + put_u64(entry + 16, layer.expert_major.experts.size); + put_u64(entry + 24, layer.expert_major.expert_stride); + put_u64(entry + 32, layer.expert_major.gate_offset); + put_u64(entry + 40, layer.expert_bytes_gate); + put_u64(entry + 48, layer.expert_major.up_offset); + put_u64(entry + 56, layer.expert_bytes_up); + put_u64(entry + 64, layer.expert_major.down_offset); + put_u64(entry + 72, layer.expert_bytes_down); + put_u64(entry + 80, layer.expert_major.gate_up_offset); + put_u64(entry + 88, layer.expert_bytes_gate_up); + } + return true; +} + +} // namespace + +uint64_t moe_expert_source_layout_hash( + const std::vector & sources, + const std::vector & layers, + const std::vector & expert_counts) { + if (layers.size() != expert_counts.size()) return 0; + uint64_t hash = 1469598103934665603ULL; + hash_u64(hash, sources.size()); + for (const MoeNvmeSource & source : sources) hash_u64(hash, source.mmap_size); + hash_u64(hash, layers.size()); + for (size_t i = 0; i < layers.size(); ++i) { + const LayerExpertRegions & layer = layers[i]; + hash_u64(hash, expert_counts[i]); + hash_u64(hash, layer.fused_gate_up ? 1 : 0); + auto add_region = [&](const ExpertFileRegion & region, + size_t expert_bytes) -> bool { + hash_u64(hash, region.source_index); + hash_u64(hash, region.offset); + hash_u64(hash, region.size); + hash_u64(hash, expert_bytes); + if (region.source_index >= sources.size()) return false; + size_t logical_bytes = 0; + if (!checked_mul(expert_bytes, (size_t) expert_counts[i], + logical_bytes) || + logical_bytes == 0 || region.size < logical_bytes) { + return false; + } + // Layout alone cannot distinguish two checkpoints with identical + // tensor shapes. Sample the start, middle and end of every expert + // stack so attaching another model's package fails closed without + // hashing hundreds of GiB at startup. + constexpr size_t kSampleBytes = 64; + uint8_t sample[kSampleBytes]{}; + const size_t sample_bytes = std::min(kSampleBytes, logical_bytes); + const size_t positions[3] = { + 0, + (logical_bytes - sample_bytes) / 2, + logical_bytes - sample_bytes, + }; + for (size_t position : positions) { + size_t offset = 0; + if (!checked_add(region.offset, position, offset) || + !read_at(sources[region.source_index], offset, + sample, sample_bytes, nullptr)) { + return false; + } + hash_u64(hash, position); + hash_bytes(hash, sample, sample_bytes); + } + return true; + }; + if (layer.fused_gate_up) { + if (!add_region(layer.gate_up_exps, + layer.expert_bytes_gate_up)) return 0; + } else { + if (!add_region(layer.gate_exps, layer.expert_bytes_gate) || + !add_region(layer.up_exps, layer.expert_bytes_up)) return 0; + } + if (!add_region(layer.down_exps, layer.expert_bytes_down)) return 0; + } + return hash; +} + +bool plan_moe_expert_package( + const std::vector & sources, + const std::vector & layers, + const std::vector & expert_counts, + const MoeExpertPackageOptions & options, + MoeExpertPackageManifest & out, + std::string * err) { + out = {}; + if (sources.empty() || layers.empty() || + layers.size() != expert_counts.size()) { + if (err) *err = "expert package source/layer dimensions are invalid"; + return false; + } + if (!is_power_of_two(options.record_alignment) || + !is_power_of_two(options.component_alignment) || + options.record_alignment < 512 || + options.component_alignment > options.record_alignment) { + if (err) *err = "expert package alignments must be powers of two"; + return false; + } + if (layers.size() > std::numeric_limits::max()) { + if (err) *err = "expert package has too many layers"; + return false; + } + for (const MoeNvmeSource & source : sources) { + if (source.mmap_size == 0 || (!source.mmap_data && source.fd < 0)) { + if (err) *err = "expert package source is unavailable"; + return false; + } + } + + size_t table_bytes = 0; + size_t raw_header_bytes = 0; + size_t cursor = 0; + if (!checked_mul(layers.size(), kLayerEntryBytes, table_bytes) || + !checked_add(kFixedHeaderBytes, table_bytes, raw_header_bytes) || + !align_up(raw_header_bytes, options.record_alignment, cursor)) { + if (err) *err = "expert package header size overflow"; + return false; + } + + MoeExpertPackageManifest manifest; + manifest.version = kVersion; + manifest.record_alignment = options.record_alignment; + manifest.component_alignment = options.component_alignment; + manifest.expert_counts = expert_counts; + manifest.layer_regions.resize(layers.size()); + + for (size_t i = 0; i < layers.size(); ++i) { + const LayerExpertRegions & source = layers[i]; + const uint32_t experts = expert_counts[i]; + if (experts == 0) { + if (err) *err = "expert package layer has zero experts"; + return false; + } + if (source.fused_gate_up) { + if (!validate_component(source.gate_up_exps, + source.expert_bytes_gate_up, experts, + sources, "gate_up", err) || + !validate_component(source.down_exps, + source.expert_bytes_down, experts, + sources, "down", err)) { + return false; + } + } else if (!validate_component(source.gate_exps, + source.expert_bytes_gate, experts, + sources, "gate", err) || + !validate_component(source.up_exps, + source.expert_bytes_up, experts, + sources, "up", err) || + !validate_component(source.down_exps, + source.expert_bytes_down, experts, + sources, "down", err)) { + return false; + } + + LayerExpertRegions packed; + packed.fused_gate_up = source.fused_gate_up; + packed.expert_bytes_gate = source.expert_bytes_gate; + packed.expert_bytes_up = source.expert_bytes_up; + packed.expert_bytes_down = source.expert_bytes_down; + packed.expert_bytes_gate_up = source.expert_bytes_gate_up; + packed.expert_major.enabled = true; + + size_t record_cursor = 0; + auto place = [&](size_t bytes, size_t & offset) -> bool { + if (!align_up(record_cursor, options.component_alignment, offset) || + !checked_add(offset, bytes, record_cursor)) { + return false; + } + return true; + }; + if (source.fused_gate_up) { + if (!place(source.expert_bytes_gate_up, + packed.expert_major.gate_up_offset) || + !place(source.expert_bytes_down, + packed.expert_major.down_offset)) { + if (err) *err = "expert package fused record size overflow"; + return false; + } + } else if (!place(source.expert_bytes_gate, + packed.expert_major.gate_offset) || + !place(source.expert_bytes_up, + packed.expert_major.up_offset) || + !place(source.expert_bytes_down, + packed.expert_major.down_offset)) { + if (err) *err = "expert package record size overflow"; + return false; + } + if (!align_up(record_cursor, options.record_alignment, + packed.expert_major.expert_stride)) { + if (err) *err = "expert package record alignment overflow"; + return false; + } + + size_t layer_bytes = 0; + if (!checked_mul(packed.expert_major.expert_stride, + (size_t) experts, layer_bytes)) { + if (err) *err = "expert package layer size overflow"; + return false; + } + packed.expert_major.experts = {cursor, layer_bytes, 0}; + if (!checked_add(cursor, layer_bytes, cursor)) { + if (err) *err = "expert package file size overflow"; + return false; + } + manifest.max_record_bytes = std::max( + manifest.max_record_bytes, packed.expert_major.expert_stride); + manifest.layer_regions[i] = packed; + } + manifest.file_bytes = cursor; + manifest.source_layout_hash = moe_expert_source_layout_hash( + sources, layers, expert_counts); + if (manifest.source_layout_hash == 0) { + if (err) *err = "failed to fingerprint expert package sources"; + return false; + } + out = std::move(manifest); + return true; +} + +bool write_moe_expert_package( + int output_fd, + const std::vector & sources, + const std::vector & layers, + const std::vector & expert_counts, + const MoeExpertPackageOptions & options, + MoeExpertPackageManifest * out, + std::string * err) { + if (output_fd < 0) { + if (err) *err = "expert package output fd is invalid"; + return false; + } + MoeExpertPackageManifest manifest; + if (!plan_moe_expert_package(sources, layers, expert_counts, + options, manifest, err)) { + return false; + } + + // Invalidate any previous package before resizing or copying data. + const uint8_t invalid_magic[8]{}; + if (!write_at(output_fd, 0, invalid_magic, sizeof(invalid_magic), err) || + !resize_file(output_fd, (size_t) manifest.file_bytes, err)) { + return false; + } + + std::vector record; + try { + record.resize(manifest.max_record_bytes); + } catch (const std::bad_alloc &) { + if (err) *err = "failed to allocate expert package record buffer"; + return false; + } + + for (size_t layer_index = 0; layer_index < layers.size(); ++layer_index) { + const LayerExpertRegions & source_layer = layers[layer_index]; + const LayerExpertRegions & packed_layer = manifest.layer_regions[layer_index]; + ComponentRef source_components[3]{}; + ComponentRef packed_components[3]{}; + const int source_count = layer_components(source_layer, source_components); + const int packed_count = layer_components(packed_layer, packed_components); + if (source_count != packed_count) { + if (err) *err = "expert package internal component mismatch"; + return false; + } + for (uint32_t expert = 0; + expert < expert_counts[layer_index]; ++expert) { + const size_t stride = packed_layer.expert_major.expert_stride; + std::memset(record.data(), 0, stride); + for (int component = 0; component < source_count; ++component) { + const ComponentRef & src = source_components[component]; + const ComponentRef & dst = packed_components[component]; + size_t expert_delta = 0; + size_t source_offset = 0; + if (!checked_mul((size_t) expert, src.bytes, expert_delta) || + !checked_add(src.source->offset, expert_delta, + source_offset) || + dst.destination_offset > stride || + dst.bytes > stride - dst.destination_offset) { + if (err) *err = "expert package component offset overflow"; + return false; + } + if (!read_at(sources[src.source->source_index], source_offset, + record.data() + dst.destination_offset, + src.bytes, err)) { + return false; + } + } + size_t record_delta = 0; + size_t output_offset = 0; + if (!checked_mul((size_t) expert, stride, record_delta) || + !checked_add(packed_layer.expert_major.experts.offset, + record_delta, output_offset) || + !write_at(output_fd, output_offset, record.data(), + stride, err)) { + return false; + } + } + if (options.progress) { + options.progress(layer_index + 1, layers.size(), + options.progress_opaque); + } + } + + if (options.sync_on_finish && !sync_file(output_fd, err)) return false; + std::vector header; + if (!encode_header(manifest, header, err) || + !write_at(output_fd, 0, header.data(), header.size(), err) || + (options.sync_on_finish && !sync_file(output_fd, err))) { + return false; + } + if (out) *out = std::move(manifest); + return true; +} + +bool read_moe_expert_package( + const MoeNvmeSource & package, + MoeExpertPackageManifest & out, + std::string * err) { + out = {}; + if (package.mmap_size < kFixedHeaderBytes || + (!package.mmap_data && package.fd < 0)) { + if (err) *err = "expert package source is unavailable or truncated"; + return false; + } + uint8_t fixed[kFixedHeaderBytes]{}; + if (!read_at(package, 0, fixed, sizeof(fixed), err)) return false; + if (std::memcmp(fixed, kMagic, sizeof(kMagic)) != 0) { + if (err) *err = "expert package magic is missing (incomplete or wrong file)"; + return false; + } + const uint32_t version = get_u32(fixed + 8); + const uint32_t header_bytes = get_u32(fixed + 12); + const uint32_t layer_count = get_u32(fixed + 16); + const uint32_t entry_bytes = get_u32(fixed + 20); + const uint64_t file_bytes = get_u64(fixed + 24); + const uint64_t source_hash = get_u64(fixed + 32); + const uint64_t record_alignment = get_u64(fixed + 40); + const uint64_t component_alignment = get_u64(fixed + 48); + if (version != kVersion || layer_count == 0 || + entry_bytes != kLayerEntryBytes || + source_hash == 0 || + file_bytes > std::numeric_limits::max() || + record_alignment > std::numeric_limits::max() || + component_alignment > std::numeric_limits::max() || + !is_power_of_two((size_t) record_alignment) || + !is_power_of_two((size_t) component_alignment) || + component_alignment > record_alignment || + header_bytes < kFixedHeaderBytes || + header_bytes > package.mmap_size || + file_bytes != package.mmap_size) { + if (err) *err = "expert package header is invalid or incompatible"; + return false; + } + size_t table_bytes = 0; + size_t required_header = 0; + if (!checked_mul((size_t) layer_count, kLayerEntryBytes, table_bytes) || + !checked_add(kFixedHeaderBytes, table_bytes, required_header) || + required_header > header_bytes) { + if (err) *err = "expert package layer table is truncated"; + return false; + } + std::vector header; + try { + header.resize(header_bytes); + } catch (const std::bad_alloc &) { + if (err) *err = "failed to allocate expert package header"; + return false; + } + if (!read_at(package, 0, header.data(), header.size(), err)) return false; + + MoeExpertPackageManifest manifest; + manifest.version = version; + manifest.source_layout_hash = source_hash; + manifest.file_bytes = file_bytes; + manifest.record_alignment = (size_t) record_alignment; + manifest.component_alignment = (size_t) component_alignment; + manifest.expert_counts.resize(layer_count); + manifest.layer_regions.resize(layer_count); + + size_t previous_end = header_bytes; + for (uint32_t i = 0; i < layer_count; ++i) { + const uint8_t * entry = header.data() + kFixedHeaderBytes + + (size_t) i * kLayerEntryBytes; + const uint32_t experts = get_u32(entry + 0); + const bool fused = (get_u32(entry + 4) & 1U) != 0; + for (size_t field = 8; field <= 88; field += 8) { + if (get_u64(entry + field) > + std::numeric_limits::max()) { + if (err) *err = "expert package layer field exceeds host size"; + return false; + } + } + const size_t data_offset = (size_t) get_u64(entry + 8); + const size_t data_bytes = (size_t) get_u64(entry + 16); + const size_t stride = (size_t) get_u64(entry + 24); + LayerExpertRegions layer; + layer.fused_gate_up = fused; + layer.expert_major.enabled = true; + layer.expert_major.experts = {data_offset, data_bytes, 0}; + layer.expert_major.expert_stride = stride; + layer.expert_major.gate_offset = (size_t) get_u64(entry + 32); + layer.expert_bytes_gate = (size_t) get_u64(entry + 40); + layer.expert_major.up_offset = (size_t) get_u64(entry + 48); + layer.expert_bytes_up = (size_t) get_u64(entry + 56); + layer.expert_major.down_offset = (size_t) get_u64(entry + 64); + layer.expert_bytes_down = (size_t) get_u64(entry + 72); + layer.expert_major.gate_up_offset = (size_t) get_u64(entry + 80); + layer.expert_bytes_gate_up = (size_t) get_u64(entry + 88); + + size_t expected_bytes = 0; + size_t data_end = 0; + if (experts == 0 || stride == 0 || + stride % record_alignment != 0 || + data_offset % record_alignment != 0 || + !checked_mul(stride, (size_t) experts, expected_bytes) || + expected_bytes != data_bytes || + data_offset < previous_end || + !checked_add(data_offset, data_bytes, data_end) || + data_end > package.mmap_size || + layer.expert_bytes_down == 0 || + (fused ? layer.expert_bytes_gate_up == 0 + : (layer.expert_bytes_gate == 0 || + layer.expert_bytes_up == 0))) { + if (err) *err = "expert package layer entry is invalid"; + return false; + } + auto contained = [&](size_t offset, size_t bytes) { + return offset <= stride && bytes <= stride - offset; + }; + if (!contained(layer.expert_major.down_offset, + layer.expert_bytes_down) || + (fused + ? !contained(layer.expert_major.gate_up_offset, + layer.expert_bytes_gate_up) + : (!contained(layer.expert_major.gate_offset, + layer.expert_bytes_gate) || + !contained(layer.expert_major.up_offset, + layer.expert_bytes_up)))) { + if (err) *err = "expert package component exceeds its record"; + return false; + } + manifest.expert_counts[i] = experts; + manifest.layer_regions[i] = layer; + manifest.max_record_bytes = std::max( + manifest.max_record_bytes, stride); + previous_end = data_end; + } + if (previous_end != package.mmap_size) { + if (err) *err = "expert package contains an unaccounted trailing range"; + return false; + } + out = std::move(manifest); + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/moe_expert_package.h b/server/src/common/moe_expert_package.h new file mode 100644 index 000000000..be8879632 --- /dev/null +++ b/server/src/common/moe_expert_package.h @@ -0,0 +1,84 @@ +// Model-neutral expert-major package for SSD-backed MoE inference. +// +// The package changes only physical byte ordering: every routed expert's +// gate/up/down components become one aligned record. Quantized bytes and the +// model's numerical contract remain unchanged. A small self-describing header +// lets any model adapter replace tensor-major shard regions with a single +// package source after validating the original layout fingerprint. + +#pragma once + +#include "moe_nvme_scheduler.h" + +#include +#include +#include +#include + +namespace dflash::common { + +struct MoeExpertPackageOptions { + // Align each complete expert record for one exact O_DIRECT request. + size_t record_alignment = 4096; + // Keep component starts suitable for backend tensor bindings/copies. + size_t component_alignment = 256; + // Publish the header only after data is durable. Disable only for disposable + // benchmark artifacts where the caller accepts crash-corrupted output. + bool sync_on_finish = true; + // Optional coarse progress hook, called after each complete layer record + // range has been written. It never runs from a background thread. + void (*progress)(size_t completed_layers, size_t total_layers, + void * opaque) = nullptr; + void * progress_opaque = nullptr; +}; + +struct MoeExpertPackageManifest { + uint32_t version = 0; + uint64_t source_layout_hash = 0; + uint64_t file_bytes = 0; + size_t record_alignment = 0; + size_t component_alignment = 0; + size_t max_record_bytes = 0; + std::vector expert_counts; + std::vector layer_regions; +}; + +// Stable fingerprint of source sizes, tensor regions, component sizes, expert +// counts, and small deterministic samples from every expert stack. It avoids +// hashing hundreds of GiB while rejecting same-shape packages from a different +// checkpoint or quantization artifact. +uint64_t moe_expert_source_layout_hash( + const std::vector & sources, + const std::vector & layers, + const std::vector & expert_counts); + +// Validate and plan a package without reading or writing weight data. +bool plan_moe_expert_package( + const std::vector & sources, + const std::vector & layers, + const std::vector & expert_counts, + const MoeExpertPackageOptions & options, + MoeExpertPackageManifest & out, + std::string * err = nullptr); + +// Compile tensor-major source regions into one expert-major output file. +// output_fd is borrowed and must be open for read/write. The file is truncated; +// its magic is published last so interrupted builds cannot look valid. +bool write_moe_expert_package( + int output_fd, + const std::vector & sources, + const std::vector & layers, + const std::vector & expert_counts, + const MoeExpertPackageOptions & options, + MoeExpertPackageManifest * out = nullptr, + std::string * err = nullptr); + +// Read and validate a package header. The returned regions all select source +// zero and can be passed directly to MoeHybridStreamEngine::bind_sources with +// the package as its only MoeNvmeSource. +bool read_moe_expert_package( + const MoeNvmeSource & package, + MoeExpertPackageManifest & out, + std::string * err = nullptr); + +} // namespace dflash::common diff --git a/server/src/common/moe_hybrid_stream.cpp b/server/src/common/moe_hybrid_stream.cpp index 72232edbf..9ae234560 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #if !defined(_WIN32) @@ -566,6 +567,7 @@ struct MoeHybridStreamEngine::Runtime { bool pending = false; bool valid = false; bool cache_managed = false; + bool pinned = false; int compute_users = 0; MoeExpertKey key{}; uint64_t frequency = 0; @@ -592,6 +594,7 @@ struct MoeHybridStreamEngine::Runtime { uint64_t device_cache_hits = 0; uint64_t device_cache_misses = 0; uint64_t device_cache_evictions = 0; + size_t pinned_experts = 0; int active_slot = -1; std::vector> graph_cache; uint64_t graph_clock = 0; @@ -615,6 +618,7 @@ void release_device_cache(RuntimeT & runtime) { } runtime.device_slots.clear(); runtime.device_index.clear(); + runtime.pinned_experts = 0; runtime.active_slot = -1; if (runtime.device_pool_buffer) { ggml_backend_buffer_free(runtime.device_pool_buffer); @@ -971,7 +975,8 @@ void MoeHybridStreamEngine::destroy() { "cache-hit=%.1f%% mean-demand-wait=%.3f ms " "dedupe=%llu upgrades=%llu dropped-prefetch=%llu " "timeouts=%llu errors=%llu " - "device-cache=%.1f MiB slots=%zu hits=%llu misses=%llu evictions=%llu " + "device-cache=%.1f MiB slots=%zu pinned=%zu " + "hits=%llu misses=%llu evictions=%llu " "graphs=%llu graph-hits=%llu graph-evictions=%llu launches=%llu\n", runtime_->io->effective_backend_name(), (unsigned long long) stats.requests, @@ -984,6 +989,7 @@ void MoeHybridStreamEngine::destroy() { (unsigned long long) stats.errors, device_cache_byte_count / 1024.0 / 1024.0, device_cache_slot_count, + runtime_->pinned_experts, (unsigned long long) runtime_->device_cache_hits, (unsigned long long) runtime_->device_cache_misses, (unsigned long long) runtime_->device_cache_evictions, @@ -1073,6 +1079,10 @@ bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, return false; } Runtime::DeviceSlot & dst = runtime_->device_slots[(size_t) device_slot]; + if (dst.pinned) { + if (err) *err = "SSD device slot is pinned by the warm expert set"; + return false; + } if (dst.compute_users != 0) { if (err) *err = "SSD device slot is still in use by expert compute"; return false; @@ -1092,6 +1102,7 @@ bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, } dst.valid = false; dst.cache_managed = false; + dst.pinned = false; dst.key = {}; dst.layout = MoeExpertIoLayout{}; dst.device_layout = Runtime::DeviceExpertLayout{}; @@ -1255,7 +1266,8 @@ bool MoeHybridStreamEngine::stage_expert_cached_async( uint64_t best_score = std::numeric_limits::max(); for (size_t i = 0; i < runtime_->device_slots.size(); ++i) { const Runtime::DeviceSlot & slot = runtime_->device_slots[i]; - if (!slot.valid || slot.pending || slot.compute_users != 0) continue; + if (!slot.valid || slot.pending || slot.pinned || + slot.compute_users != 0) continue; const uint64_t age = runtime_->device_clock >= slot.last_touch ? runtime_->device_clock - slot.last_touch : 0; const uint64_t recency = age < 65535 ? 65535 - age : 0; @@ -1343,6 +1355,109 @@ size_t MoeHybridStreamEngine::device_cache_bytes() const { return runtime_ ? runtime_->device_pool_bytes : 0; } +size_t MoeHybridStreamEngine::pinned_expert_count() const { + return runtime_ ? runtime_->pinned_experts : 0; +} + +bool MoeHybridStreamEngine::warm_and_pin_device_cache( + const std::vector & layer_specs, + const std::vector & entries, + int reserve_slots, + MoeStreamCacheWarmStats * stats, + std::string * err) { + MoeStreamCacheWarmStats local; + local.requested = entries.size(); + if (stats) *stats = local; + if (!is_bound()) { + if (err) *err = "cannot warm an unbound streamed expert cache"; + return false; + } + reserve_slots = std::max(2, reserve_slots); + + std::lock_guard compute_guard(runtime_->compute_mutex); + ScopedGpuDevice device_scope(runtime_->device); + if (!device_scope.ready()) { + if (err) *err = "failed to select streamed cache GPU for warmup"; + return false; + } + + std::vector candidates = entries; + std::stable_sort(candidates.begin(), candidates.end(), + [](const MoeStreamCacheWarmEntry & a, + const MoeStreamCacheWarmEntry & b) { + if (a.frequency != b.frequency) return a.frequency > b.frequency; + if (a.layer != b.layer) return a.layer < b.layer; + return a.expert < b.expert; + }); + std::unordered_set seen; + std::vector unique; + unique.reserve(candidates.size()); + for (const MoeStreamCacheWarmEntry & candidate : candidates) { + if (candidate.layer < 0 || candidate.expert < 0 || + (size_t) candidate.layer >= layer_specs.size()) { + if (err) *err = "streamed cache warm entry is out of range"; + return false; + } + const uint64_t key = device_key(candidate.layer, candidate.expert); + if (seen.insert(key).second) unique.push_back(candidate); + } + + // Establish every backend-padded layer layout before loading data. If one + // format needs a larger device stride, the cache is resized once while it + // is still empty instead of invalidating an already-warmed prefix. + std::unordered_set prepared_layers; + for (const MoeStreamCacheWarmEntry & candidate : unique) { + if (!prepared_layers.insert(candidate.layer).second) continue; + if (!prepare_device_expert_layout( + *runtime_, candidate.layer, + layer_specs[(size_t) candidate.layer], err)) { + return false; + } + } + + const size_t slot_count = runtime_->device_slots.size(); + const size_t pin_capacity = slot_count > (size_t) reserve_slots + ? slot_count - (size_t) reserve_slots : 0; + for (const MoeStreamCacheWarmEntry & candidate : unique) { + const uint64_t key = device_key(candidate.layer, candidate.expert); + auto found = runtime_->device_index.find(key); + if (found != runtime_->device_index.end()) { + const int slot_index = found->second; + if (slot_index >= 0 && slot_index < (int) slot_count) { + Runtime::DeviceSlot & resident = + runtime_->device_slots[(size_t) slot_index]; + if (resident.valid && resident.pinned) { + ++local.already_resident; + continue; + } + } + } + if (runtime_->pinned_experts >= pin_capacity) { + ++local.capacity_drops; + continue; + } + + const bool was_resident = found != runtime_->device_index.end(); + int device_slot = -1; + if (!stage_expert_cached_async( + candidate.layer, candidate.expert, &device_slot, err) || + !activate_device_slot(device_slot, err)) { + if (stats) *stats = local; + return false; + } + release_device_slot(device_slot); + Runtime::DeviceSlot & slot = + runtime_->device_slots[(size_t) device_slot]; + slot.pinned = true; + slot.frequency = std::max(slot.frequency, 2); + ++runtime_->pinned_experts; + ++local.admitted; + if (was_resident) ++local.already_resident; + } + if (stats) *stats = local; + return true; +} + ggml_backend_t MoeHybridStreamEngine::compute_backend() const { return runtime_ ? runtime_->backend : nullptr; } @@ -2002,6 +2117,22 @@ uint32_t stream_owner_hash(int layer, int expert) { } // namespace +bool moe_stream_primary_owns_expert( + const MoeStreamDualOwnerPolicy & policy, + int layer, + int expert) { + if (layer < 0 || expert < 0 || + policy.primary_share_per_mille < 0 || + policy.primary_share_per_mille > 1000) { + return false; + } + if (policy.primary_placement) { + return policy.primary_placement->is_hot(layer, expert); + } + return stream_owner_hash(layer, expert) % 1000U < + (uint32_t) policy.primary_share_per_mille; +} + bool partition_moe_stream_routes( const MoeStreamRouteBatch & batch, const MoeStreamDualOwnerPolicy & policy, @@ -2049,15 +2180,8 @@ bool partition_moe_stream_routes( return false; } if (owner[(size_t) expert] >= 0) continue; - bool primary_owner = false; - if (policy.primary_placement) { - primary_owner = policy.primary_placement->is_hot( - batch.layer, expert); - } else { - primary_owner = - stream_owner_hash(batch.layer, expert) % 1000U < - (uint32_t) policy.primary_share_per_mille; - } + const bool primary_owner = moe_stream_primary_owns_expert( + policy, batch.layer, expert); owner[(size_t) expert] = primary_owner ? 1 : 0; unique_experts.push_back(expert); if (primary_owner) ++primary_experts; diff --git a/server/src/common/moe_hybrid_stream.h b/server/src/common/moe_hybrid_stream.h index 1c7ed1c3b..9eee5169c 100644 --- a/server/src/common/moe_hybrid_stream.h +++ b/server/src/common/moe_hybrid_stream.h @@ -102,6 +102,20 @@ struct MoeStreamComputeStats { uint64_t graph_launches = 0; }; +struct MoeStreamCacheWarmEntry { + int32_t layer = -1; + int32_t expert = -1; + uint64_t frequency = 0; + uint64_t bytes = 0; +}; + +struct MoeStreamCacheWarmStats { + size_t requested = 0; + size_t admitted = 0; + size_t already_resident = 0; + size_t capacity_drops = 0; +}; + // Route ownership for two concurrent SSD-backed GPU owners. An explicit // placement takes precedence and identifies the primary GPU's hot experts. // Without one, a stable layer/expert hash supplies a deterministic capacity @@ -183,8 +197,19 @@ class MoeHybridStreamEngine { void release_device_slot(int device_slot); int device_slot_count() const; size_t device_cache_bytes() const; + size_t pinned_expert_count() const; ggml_backend_t compute_backend() const; + // Populate and protect the highest-value profile entries. Numerical specs + // are supplied per layer so mixed-format models remain valid. At least + // reserve_slots stay evictable for the ordinary miss pipeline. + bool warm_and_pin_device_cache( + const std::vector & layer_specs, + const std::vector & entries, + int reserve_slots, + MoeStreamCacheWarmStats * stats = nullptr, + std::string * err = nullptr); + bool stream_expert_sync(int layer, int expert_id, std::string * err = nullptr); @@ -281,6 +306,12 @@ bool partition_moe_stream_routes( MoeStreamDualOwnerStats * stats = nullptr, std::string * err = nullptr); +// Stable owner decision shared by route partitioning and offline cache plans. +bool moe_stream_primary_owns_expert( + const MoeStreamDualOwnerPolicy & policy, + int layer, + int expert); + // Evaluate the cold contribution for one layer. All routed SSD requests are // admitted before compute starts, then double-buffered H2D runs concurrently // with the preceding expert graph. diff --git a/server/src/common/moe_stream_cache_policy.cpp b/server/src/common/moe_stream_cache_policy.cpp new file mode 100644 index 000000000..fb5b69a2d --- /dev/null +++ b/server/src/common/moe_stream_cache_policy.cpp @@ -0,0 +1,91 @@ +#include "moe_stream_cache_policy.h" + +#include +#include +#include + +namespace dflash::common { + +bool build_moe_stream_cache_plan( + const MoeHybridRoutingStats & stats, + const std::vector & layer_expert_bytes, + const MoeStreamCachePlanConfig & config, + const MoeStreamDualOwnerPolicy * owner_policy, + std::vector & out, + std::string * err) { + out.clear(); + if (stats.empty() || stats.n_layer <= 0 || stats.n_expert <= 0 || + layer_expert_bytes.size() != (size_t) stats.n_layer) { + if (err) *err = "invalid routing profile or per-layer expert sizes"; + return false; + } + if (config.max_entries == 0) return true; + if (config.owner != MoeStreamCacheOwner::All) { + if (!owner_policy || owner_policy->primary_share_per_mille < 0 || + owner_policy->primary_share_per_mille > 1000 || + (owner_policy->primary_placement && + (owner_policy->primary_placement->n_layer != stats.n_layer || + owner_policy->primary_placement->n_expert != stats.n_expert))) { + if (err) *err = "cache owner policy does not match routing profile"; + return false; + } + } + + std::vector candidates; + try { + candidates.reserve((size_t) stats.n_layer * (size_t) stats.n_expert); + } catch (const std::bad_alloc &) { + if (err) *err = "failed to allocate streamed cache candidates"; + return false; + } + for (int layer = 0; layer < stats.n_layer; ++layer) { + const uint64_t bytes = layer_expert_bytes[(size_t) layer]; + if (bytes == 0) { + if (err) *err = "streamed cache plan has a zero-sized expert"; + return false; + } + for (int expert = 0; expert < stats.n_expert; ++expert) { + const uint64_t frequency = stats.count(layer, expert); + if (frequency == 0) continue; + if (config.owner != MoeStreamCacheOwner::All) { + const bool primary = moe_stream_primary_owns_expert( + *owner_policy, layer, expert); + if ((config.owner == MoeStreamCacheOwner::Primary) != primary) { + continue; + } + } + candidates.push_back({ + (int32_t) layer, (int32_t) expert, frequency, bytes}); + } + } + + std::stable_sort(candidates.begin(), candidates.end(), + [](const MoeStreamCacheWarmEntry & a, + const MoeStreamCacheWarmEntry & b) { + const long double av = (long double) a.frequency / + (long double) a.bytes; + const long double bv = (long double) b.frequency / + (long double) b.bytes; + if (av != bv) return av > bv; + if (a.frequency != b.frequency) return a.frequency > b.frequency; + if (a.layer != b.layer) return a.layer < b.layer; + return a.expert < b.expert; + }); + + uint64_t used_bytes = 0; + out.reserve(std::min(config.max_entries, candidates.size())); + for (const MoeStreamCacheWarmEntry & candidate : candidates) { + if (out.size() >= config.max_entries) break; + if (config.max_bytes != 0) { + if (candidate.bytes > config.max_bytes - + std::min(config.max_bytes, used_bytes)) { + continue; + } + used_bytes += candidate.bytes; + } + out.push_back(candidate); + } + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/moe_stream_cache_policy.h b/server/src/common/moe_stream_cache_policy.h new file mode 100644 index 000000000..21ef0a1ed --- /dev/null +++ b/server/src/common/moe_stream_cache_policy.h @@ -0,0 +1,42 @@ +// Profile-guided, model-neutral warm-cache planning for streamed MoE experts. + +#pragma once + +#include "moe_hybrid_routing_stats.h" +#include "moe_hybrid_stream.h" + +#include +#include +#include +#include + +namespace dflash::common { + +enum class MoeStreamCacheOwner { + All, + Primary, + Secondary, +}; + +struct MoeStreamCachePlanConfig { + size_t max_entries = 0; + // Zero means entry-count budgeting only. A nonzero budget makes the + // planner rank by observed frequency per byte and skip entries that do not + // fit the remaining capacity. + uint64_t max_bytes = 0; + MoeStreamCacheOwner owner = MoeStreamCacheOwner::All; +}; + +// Select a deterministic highest-value warm set. layer_expert_bytes contains +// one complete routed expert size per layer. Owner filtering reuses the exact +// runtime partition policy, so warming cannot duplicate an expert across the +// two Lucebox GPUs. +bool build_moe_stream_cache_plan( + const MoeHybridRoutingStats & stats, + const std::vector & layer_expert_bytes, + const MoeStreamCachePlanConfig & config, + const MoeStreamDualOwnerPolicy * owner_policy, + std::vector & out, + std::string * err = nullptr); + +} // namespace dflash::common diff --git a/server/src/kimi_k3/kimi_k3_backend.cpp b/server/src/kimi_k3/kimi_k3_backend.cpp index 6395086b2..eb450a56f 100644 --- a/server/src/kimi_k3/kimi_k3_backend.cpp +++ b/server/src/kimi_k3/kimi_k3_backend.cpp @@ -1,6 +1,8 @@ #include "kimi_k3_backend.h" +#include "common/moe_expert_package.h" #include "common/moe_hybrid_placement.h" +#include "common/moe_stream_cache_policy.h" #include "common/sampler.h" #include "dflash27b.h" @@ -14,12 +16,15 @@ #include #include #include +#include #include #if defined(_WIN32) #include #include +#include #include +#include #else #include #include @@ -27,6 +32,129 @@ #endif namespace dflash::common { +namespace { + +void close_file_descriptor(int fd) { +#if defined(_WIN32) + ::_close(fd); +#else + ::close(fd); +#endif +} + +class ScopedFileDescriptors { +public: + ~ScopedFileDescriptors() { + for (int fd : values_) close_file_descriptor(fd); + } + + void add(int fd) { values_.push_back(fd); } + + void close(int fd) { + const auto found = std::find(values_.begin(), values_.end(), fd); + if (found == values_.end()) return; + close_file_descriptor(fd); + values_.erase(found); + } + +private: + std::vector values_; +}; + +bool open_nvme_source(const std::string & path, + ScopedFileDescriptors & descriptors, + MoeNvmeSource & out, + std::string & error) { +#if defined(_WIN32) + const int fd = ::_open(path.c_str(), _O_RDONLY | _O_BINARY); +#else + const int fd = ::open(path.c_str(), O_RDONLY | O_CLOEXEC); +#endif + if (fd < 0) { + error = "cannot open " + path + ": " + std::strerror(errno); + return false; + } + uint64_t bytes = 0; +#if defined(_WIN32) + struct _stat64 stat_buffer {}; + if (::_fstat64(fd, &stat_buffer) == 0 && stat_buffer.st_size > 0) { + bytes = static_cast(stat_buffer.st_size); + } +#else + struct stat stat_buffer {}; + if (::fstat(fd, &stat_buffer) == 0 && stat_buffer.st_size > 0) { + bytes = static_cast(stat_buffer.st_size); + } +#endif + if (bytes == 0 || bytes > std::numeric_limits::max()) { + close_file_descriptor(fd); + error = "cannot determine file size for " + path; + return false; + } + descriptors.add(fd); + out = {nullptr, static_cast(bytes), fd}; + return true; +} + +bool file_path_exists(const char * path) { + if (!path || !*path) return false; +#if defined(_WIN32) + struct _stat64 st {}; + return ::_stat64(path, &st) == 0 && st.st_size > 0; +#else + struct stat st {}; + return ::stat(path, &st) == 0 && st.st_size > 0; +#endif +} + +bool create_package_output(const std::string & path, + ScopedFileDescriptors & descriptors, + int & fd, std::string & error) { +#if defined(_WIN32) + fd = ::_open(path.c_str(), + _O_CREAT | _O_TRUNC | _O_RDWR | _O_BINARY, + _S_IREAD | _S_IWRITE); +#else + fd = ::open(path.c_str(), + O_CREAT | O_TRUNC | O_RDWR | O_CLOEXEC, 0644); +#endif + if (fd < 0) { + error = "cannot create " + path + ": " + std::strerror(errno); + return false; + } + descriptors.add(fd); + return true; +} + +std::string package_temporary_path(const std::string & path) { +#if defined(_WIN32) + const int process_id = ::_getpid(); +#else + const int process_id = static_cast(::getpid()); +#endif + return path + ".partial." + std::to_string(process_id); +} + +bool publish_package(const std::string & temporary, + const std::string & destination, + std::string & error) { +#if defined(_WIN32) + if (::MoveFileExA(temporary.c_str(), destination.c_str(), + MOVEFILE_REPLACE_EXISTING | + MOVEFILE_WRITE_THROUGH)) { + return true; + } + error = "cannot publish expert package " + destination + + ": Windows error " + std::to_string(::GetLastError()); +#else + if (::rename(temporary.c_str(), destination.c_str()) == 0) return true; + error = "cannot publish expert package " + destination + ": " + + std::strerror(errno); +#endif + return false; +} + +} // namespace KimiK3Backend::KimiK3Backend(const KimiK3BackendConfig & cfg) : cfg_(cfg) {} @@ -43,6 +171,8 @@ void KimiK3Backend::release_expert_backend() { } bool KimiK3Backend::init_streaming() { + routing_stats_.reset(); + routing_stats_out_path_.clear(); if (!weights_.routed_experts_streamed || weights_.streamed_layer_regions.empty() || weights_.max_streamed_expert_bytes == 0) { @@ -137,119 +267,200 @@ bool KimiK3Backend::init_streaming() { return stream_config; }; - const MoeStreamConfig primary_config = stream_config_for( - cfg_.device.primary_gpu(), "primary"); - if (!stream_engine_.init( - backend_, weights_.max_streamed_expert_bytes, - primary_config, &error)) { - std::fprintf(stderr, - "[kimi-k3] primary stream engine initialization failed: %s\n", - error.c_str()); - return fail_streaming(); - } - if (expert_backend_) { - const MoeStreamConfig secondary_config = stream_config_for( - expert_gpu_, "secondary"); - if (!secondary_stream_engine_.init( - expert_backend_, weights_.max_streamed_expert_bytes, - secondary_config, &error)) { - std::fprintf(stderr, - "[kimi-k3] secondary stream engine initialization failed: %s\n", - error.c_str()); - return fail_streaming(); - } - } - stream_owner_policy_ = MoeStreamDualOwnerPolicy::from_env(); stream_placement_ = MoeHybridPlacement{}; const char * placement_path = std::getenv("DFLASH_MOE_PLACEMENT"); - if (expert_backend_ && placement_path && *placement_path) { + if (placement_path && *placement_path) { if (!MoeHybridPlacement::load_json( placement_path, stream_placement_, &error) || !stream_placement_.matches( static_cast(weights_.streamed_layer_regions.size()), weights_.n_expert, weights_.n_expert_used)) { std::fprintf(stderr, - "[kimi-k3] invalid dual-owner placement %s: %s\n", + "[kimi-k3] invalid streamed-expert placement %s: %s\n", placement_path, error.empty() ? "model shape mismatch" : error.c_str()); return fail_streaming(); } - stream_owner_policy_.primary_placement = &stream_placement_; + if (expert_backend_) { + stream_owner_policy_.primary_placement = &stream_placement_; + } } - std::vector descriptors; - std::vector sources; - descriptors.reserve(weights_.shard_paths.size()); - sources.reserve(weights_.shard_paths.size()); + ScopedFileDescriptors descriptors; + std::vector model_sources; + model_sources.reserve(weights_.shard_paths.size()); for (const std::string & shard : weights_.shard_paths) { -#if defined(_WIN32) - const int fd = ::_open(shard.c_str(), _O_RDONLY | _O_BINARY); -#else - const int fd = ::open(shard.c_str(), O_RDONLY | O_CLOEXEC); -#endif - if (fd < 0) { + MoeNvmeSource source; + if (!open_nvme_source(shard, descriptors, source, error)) { std::fprintf(stderr, - "[kimi-k3] cannot open expert shard %s: %s\n", - shard.c_str(), std::strerror(errno)); - for (int opened : descriptors) { -#if defined(_WIN32) - ::_close(opened); -#else - ::close(opened); -#endif - } + "[kimi-k3] %s\n", error.c_str()); return fail_streaming(); } - uint64_t shard_bytes = 0; -#if defined(_WIN32) - struct _stat64 stat_buffer {}; - if (::_fstat64(fd, &stat_buffer) == 0 && stat_buffer.st_size > 0) { - shard_bytes = static_cast(stat_buffer.st_size); - } -#else - struct stat stat_buffer {}; - if (::fstat(fd, &stat_buffer) == 0 && stat_buffer.st_size > 0) { - shard_bytes = static_cast(stat_buffer.st_size); + model_sources.push_back(source); + } + + std::vector expert_counts( + weights_.streamed_layer_regions.size(), + static_cast(weights_.n_expert)); + const uint64_t source_layout_hash = moe_expert_source_layout_hash( + model_sources, weights_.streamed_layer_regions, expert_counts); + std::vector package_sources; + MoeExpertPackageManifest package_manifest; + const std::vector * active_sources = &model_sources; + const std::vector * active_regions = + &weights_.streamed_layer_regions; + size_t max_streamed_expert_bytes = weights_.max_streamed_expert_bytes; + const char * package_path = std::getenv("DFLASH_MOE_EXPERT_PACKAGE"); + const char * package_build = + std::getenv("DFLASH_MOE_EXPERT_PACKAGE_BUILD"); + const bool build_requested = + package_build && *package_build && + std::strcmp(package_build, "0") != 0; + const bool force_build = + build_requested && std::strcmp(package_build, "force") == 0; + if (build_requested && (!package_path || !*package_path)) { + std::fprintf(stderr, + "[kimi-k3] DFLASH_MOE_EXPERT_PACKAGE_BUILD requires " + "DFLASH_MOE_EXPERT_PACKAGE=\n"); + return fail_streaming(); + } + if (package_path && *package_path) { + MoeNvmeSource package_source; + auto load_matching_package = [&]() -> bool { + MoeNvmeSource candidate; + MoeExpertPackageManifest candidate_manifest; + if (!open_nvme_source( + package_path, descriptors, candidate, error)) { + return false; + } + if (!read_moe_expert_package( + candidate, candidate_manifest, &error)) { + descriptors.close(candidate.fd); + return false; + } + const bool shape_matches = + candidate_manifest.layer_regions.size() == + weights_.streamed_layer_regions.size() && + candidate_manifest.expert_counts == expert_counts; + if (!shape_matches || + candidate_manifest.source_layout_hash != source_layout_hash) { + descriptors.close(candidate.fd); + error = "expert package does not match this model, " + "checkpoint, or shard layout"; + return false; + } + package_source = candidate; + package_manifest = std::move(candidate_manifest); + return true; + }; + + bool package_loaded = !force_build && load_matching_package(); + if (!package_loaded && !build_requested) { + std::fprintf(stderr, + "[kimi-k3] invalid expert package %s: %s\n", + package_path, error.c_str()); + return fail_streaming(); } -#endif - if (shard_bytes == 0 || - shard_bytes > std::numeric_limits::max()) { + if (!package_loaded) { + if (!force_build && file_path_exists(package_path)) { + std::fprintf(stderr, + "[kimi-k3] rebuilding invalid expert package %s: %s\n", + package_path, error.c_str()); + } + error.clear(); + const std::string temporary_path = + package_temporary_path(package_path); + int output_fd = -1; + if (!create_package_output( + temporary_path, descriptors, output_fd, error)) { + std::fprintf(stderr, + "[kimi-k3] expert package build failed: %s\n", + error.c_str()); + return fail_streaming(); + } + MoeExpertPackageOptions options; + options.progress = [](size_t completed, size_t total, void *) { + std::fprintf(stderr, + "[kimi-k3] expert package progress %zu/%zu layers\n", + completed, total); + std::fflush(stderr); + }; std::fprintf(stderr, - "[kimi-k3] cannot determine expert shard size: %s\n", - shard.c_str()); -#if defined(_WIN32) - ::_close(fd); -#else - ::close(fd); -#endif - for (int opened : descriptors) { -#if defined(_WIN32) - ::_close(opened); -#else - ::close(opened); -#endif + "[kimi-k3] compiling expert-major package=%s " + "(exact weights, one aligned record/expert)\n", + package_path); + if (!write_moe_expert_package( + output_fd, model_sources, + weights_.streamed_layer_regions, expert_counts, + options, &package_manifest, &error)) { + descriptors.close(output_fd); + (void) std::remove(temporary_path.c_str()); + std::fprintf(stderr, + "[kimi-k3] expert package build failed: %s\n", + error.c_str()); + return fail_streaming(); + } + descriptors.close(output_fd); + if (!publish_package( + temporary_path, package_path, error)) { + (void) std::remove(temporary_path.c_str()); + std::fprintf(stderr, + "[kimi-k3] expert package build failed: %s\n", + error.c_str()); + return fail_streaming(); } + package_loaded = load_matching_package(); + if (!package_loaded) { + std::fprintf(stderr, + "[kimi-k3] published expert package is invalid %s: %s\n", + package_path, error.c_str()); + return fail_streaming(); + } + } + package_sources.push_back(package_source); + active_sources = &package_sources; + active_regions = &package_manifest.layer_regions; + max_streamed_expert_bytes = std::max( + max_streamed_expert_bytes, package_manifest.max_record_bytes); + std::fprintf(stderr, + "[kimi-k3] using expert-major package=%s layers=%zu " + "record<=%.2f MiB (one read/expert)\n", + package_path, package_manifest.layer_regions.size(), + static_cast(package_manifest.max_record_bytes) / + (1024.0 * 1024.0)); + } + + const MoeStreamConfig primary_config = stream_config_for( + cfg_.device.primary_gpu(), "primary"); + if (!stream_engine_.init( + backend_, max_streamed_expert_bytes, + primary_config, &error)) { + std::fprintf(stderr, + "[kimi-k3] primary stream engine initialization failed: %s\n", + error.c_str()); + return fail_streaming(); + } + if (expert_backend_) { + const MoeStreamConfig secondary_config = stream_config_for( + expert_gpu_, "secondary"); + if (!secondary_stream_engine_.init( + expert_backend_, max_streamed_expert_bytes, + secondary_config, &error)) { + std::fprintf(stderr, + "[kimi-k3] secondary stream engine initialization failed: %s\n", + error.c_str()); return fail_streaming(); } - descriptors.push_back(fd); - sources.push_back({ - nullptr, static_cast(shard_bytes), fd}); } + const bool primary_bound = stream_engine_.bind_sources( - sources, weights_.streamed_layer_regions, &error); + *active_sources, *active_regions, &error); bool secondary_bound = true; std::string secondary_error; if (primary_bound && expert_backend_) { secondary_bound = secondary_stream_engine_.bind_sources( - sources, weights_.streamed_layer_regions, &secondary_error); - } - for (int fd : descriptors) { -#if defined(_WIN32) - ::_close(fd); -#else - ::close(fd); -#endif + *active_sources, *active_regions, &secondary_error); } if (!primary_bound || !secondary_bound) { std::fprintf(stderr, @@ -257,6 +468,134 @@ bool KimiK3Backend::init_streaming() { primary_bound ? secondary_error.c_str() : error.c_str()); return fail_streaming(); } + + std::vector layer_specs; + std::vector layer_expert_bytes; + layer_specs.reserve(weights_.streamed_layer_regions.size()); + layer_expert_bytes.reserve(weights_.streamed_layer_regions.size()); + for (size_t local_layer = 0; + local_layer < weights_.streamed_layer_regions.size(); ++local_layer) { + const size_t model_layer = + (size_t) weights_.n_dense_lead + local_layer; + if (model_layer >= weights_.layers.size()) { + std::fprintf(stderr, + "[kimi-k3] streamed layer metadata is out of range\n"); + return fail_streaming(); + } + const KimiK3Layer & layer = weights_.layers[model_layer]; + if (!layer.ffn_gate_exps || !layer.ffn_up_exps || + !layer.ffn_down_exps) { + std::fprintf(stderr, + "[kimi-k3] streamed layer is missing expert types\n"); + return fail_streaming(); + } + MoeStreamExpertSpec spec; + spec.input_dim = weights_.n_expert_latent; + spec.intermediate_dim = weights_.n_ff_exp; + spec.output_dim = weights_.n_expert_latent; + spec.gate_type = layer.ffn_gate_exps->type; + spec.up_type = layer.ffn_up_exps->type; + spec.down_type = layer.ffn_down_exps->type; + spec.gated_activation = MoeGatedActivation::Situ; + spec.situ_beta = weights_.situ_beta; + spec.situ_linear_beta = weights_.situ_linear_beta; + layer_specs.push_back(spec); + const LayerExpertRegions & regions = + weights_.streamed_layer_regions[local_layer]; + layer_expert_bytes.push_back( + (uint64_t) regions.expert_bytes_gate + + (uint64_t) regions.expert_bytes_up + + (uint64_t) regions.expert_bytes_down + + (uint64_t) regions.expert_bytes_gate_up); + } + + const char * hotness_path = std::getenv("DFLASH_MOE_HOTNESS_CSV"); + if (!hotness_path || !*hotness_path) { + hotness_path = std::getenv("DFLASH_DS4_HOTNESS_CSV"); + } + MoeHybridRoutingStats routing_profile; + const bool have_routing_profile = hotness_path && *hotness_path; + if (have_routing_profile && + (!MoeHybridRoutingStats::load_csv( + hotness_path, routing_profile, &error) || + !routing_profile.matches( + static_cast(weights_.streamed_layer_regions.size()), + weights_.n_expert, weights_.n_expert_used))) { + std::fprintf(stderr, + "[kimi-k3] invalid expert hotness profile %s: %s\n", + hotness_path, + error.empty() ? "model shape mismatch" : error.c_str()); + return fail_streaming(); + } + + auto warm_owner = [&](MoeHybridStreamEngine & engine, + MoeStreamCacheOwner owner_kind, + const char * owner_name) -> bool { + const int slots = engine.device_slot_count(); + // Preserve a meaningful adaptive region when the profile drifts. Tiny + // caches still retain the two slots needed by the miss pipeline. + const int reserve_slots = std::max(2, slots / 4); + const size_t max_entries = slots > reserve_slots + ? (size_t) (slots - reserve_slots) : 0; + if (max_entries == 0) return true; + + std::vector plan; + if (have_routing_profile) { + MoeStreamCachePlanConfig plan_config; + plan_config.max_entries = max_entries; + plan_config.owner = owner_kind; + const MoeStreamDualOwnerPolicy * policy = + owner_kind == MoeStreamCacheOwner::All + ? nullptr : &stream_owner_policy_; + if (!build_moe_stream_cache_plan( + routing_profile, layer_expert_bytes, + plan_config, policy, plan, &error)) { + return false; + } + } else if (!stream_placement_.empty() && + owner_kind != MoeStreamCacheOwner::Secondary) { + for (int layer = 0; layer < stream_placement_.n_layer; ++layer) { + const auto & ids = + stream_placement_.hot_expert_ids[(size_t) layer]; + for (size_t rank = 0; rank < ids.size(); ++rank) { + plan.push_back({ + (int32_t) layer, ids[rank], + (uint64_t) std::max(1, ids.size() - rank), + layer_expert_bytes[(size_t) layer]}); + } + } + } + if (plan.empty()) return true; + + MoeStreamCacheWarmStats warm_stats; + if (!engine.warm_and_pin_device_cache( + layer_specs, plan, reserve_slots, &warm_stats, &error)) { + return false; + } + std::fprintf(stderr, + "[kimi-k3] %s profile-warm cache: requested=%zu pinned=%zu " + "resident=%zu capacity-drops=%zu source=%s\n", + owner_name, warm_stats.requested, warm_stats.admitted, + warm_stats.already_resident, warm_stats.capacity_drops, + have_routing_profile ? hotness_path : placement_path); + return true; + }; + + if (!warm_owner( + stream_engine_, + expert_backend_ ? MoeStreamCacheOwner::Primary + : MoeStreamCacheOwner::All, + "primary") || + (expert_backend_ && + !warm_owner(secondary_stream_engine_, + MoeStreamCacheOwner::Secondary, + "secondary"))) { + std::fprintf(stderr, + "[kimi-k3] profile-guided cache warmup failed: %s\n", + error.c_str()); + return fail_streaming(); + } + if (expert_backend_ && !dual_stream_executor_.init( stream_engine_, secondary_stream_engine_, &error)) { std::fprintf(stderr, @@ -264,6 +603,22 @@ bool KimiK3Backend::init_streaming() { error.c_str()); return fail_streaming(); } + if (const char * stats_path = std::getenv("DFLASH_MOE_ROUTE_STATS_OUT")) { + if (*stats_path) { + routing_stats_ = std::make_shared(); + if (!routing_stats_->init( + static_cast(weights_.streamed_layer_regions.size()), + weights_.n_expert, weights_.n_expert_used)) { + std::fprintf(stderr, + "[kimi-k3] failed to initialize routed-expert statistics\n"); + return fail_streaming(); + } + routing_stats_out_path_ = stats_path; + std::fprintf(stderr, + "[kimi-k3] recording native route counts to %s\n", + routing_stats_out_path_.c_str()); + } + } if (expert_backend_) { std::fprintf(stderr, "[kimi-k3] routed experts dual-owner: shards=%zu layers=%zu " @@ -341,6 +696,7 @@ void KimiK3Backend::print_ready_banner() const { bool KimiK3Backend::park(ParkTarget target) { if (!park_target_includes_target_model(target)) return false; if (!parked_) { + maybe_save_routing_stats(); dual_stream_executor_.destroy(); stream_engine_.destroy(); secondary_stream_engine_.destroy(); @@ -408,7 +764,7 @@ GenerateResult KimiK3Backend::generate_impl(const GenerateRequest & req, static_cast(i), logits, &stream_engine_, dual_stream_executor_.is_ready() ? &dual_stream_executor_ : nullptr, - &stream_owner_policy_)) { + &stream_owner_policy_, routing_stats_.get())) { result.fail(GenerateErrorCode::PrefillFailed, dflash27b_last_error()); out_io.emit(-1); @@ -454,7 +810,7 @@ GenerateResult KimiK3Backend::generate_impl(const GenerateRequest & req, cache_.cur_pos, logits, &stream_engine_, dual_stream_executor_.is_ready() ? &dual_stream_executor_ : nullptr, - &stream_owner_policy_)) { + &stream_owner_policy_, routing_stats_.get())) { result.fail(GenerateErrorCode::DecodeFailed, dflash27b_last_error()); out_io.emit(-1); @@ -464,6 +820,7 @@ GenerateResult KimiK3Backend::generate_impl(const GenerateRequest & req, } const auto decode_end = std::chrono::steady_clock::now(); result.decode_s = std::chrono::duration(decode_end - decode_begin).count(); + maybe_save_routing_stats(); out_io.emit(-1); result.succeed(); return result; @@ -507,6 +864,7 @@ bool KimiK3Backend::handle_compress(const std::string & line, } void KimiK3Backend::shutdown() { + maybe_save_routing_stats(); dual_stream_executor_.destroy(); stream_engine_.destroy(); secondary_stream_engine_.destroy(); @@ -517,7 +875,19 @@ void KimiK3Backend::shutdown() { ggml_backend_free(backend_); backend_ = nullptr; } + routing_stats_.reset(); + routing_stats_out_path_.clear(); parked_ = false; } +void KimiK3Backend::maybe_save_routing_stats() { + if (!routing_stats_ || routing_stats_out_path_.empty()) return; + std::string error; + if (!routing_stats_->save_csv(routing_stats_out_path_, &error)) { + std::fprintf(stderr, + "[kimi-k3] failed to save route statistics %s: %s\n", + routing_stats_out_path_.c_str(), error.c_str()); + } +} + } // namespace dflash::common diff --git a/server/src/kimi_k3/kimi_k3_backend.h b/server/src/kimi_k3/kimi_k3_backend.h index 64e7453ef..ac01d0065 100644 --- a/server/src/kimi_k3/kimi_k3_backend.h +++ b/server/src/kimi_k3/kimi_k3_backend.h @@ -1,12 +1,14 @@ #pragma once #include "common/model_backend.h" +#include "common/moe_hybrid_routing_stats.h" #include "common/moe_hybrid_stream.h" #include "common/moe_storage_policy.h" #include "kimi_k3_internal.h" #include "placement/placement_config.h" #include +#include #include namespace dflash::common { @@ -56,6 +58,7 @@ class KimiK3Backend final : public ModelBackend { private: bool init_streaming(); void release_expert_backend(); + void maybe_save_routing_stats(); int32_t choose_token(const std::vector & logits, const SamplerCfg & sampler, @@ -72,6 +75,8 @@ class KimiK3Backend final : public ModelBackend { MoeStreamDualOwnerExecutor dual_stream_executor_; MoeHybridPlacement stream_placement_; MoeStreamDualOwnerPolicy stream_owner_policy_; + std::shared_ptr routing_stats_; + std::string routing_stats_out_path_; bool parked_ = false; std::mt19937_64 rng_{std::random_device{}()}; }; diff --git a/server/src/kimi_k3/kimi_k3_graph.cpp b/server/src/kimi_k3/kimi_k3_graph.cpp index cc78dfd27..38788a4b4 100644 --- a/server/src/kimi_k3/kimi_k3_graph.cpp +++ b/server/src/kimi_k3/kimi_k3_graph.cpp @@ -1,5 +1,6 @@ #include "kimi_k3_internal.h" +#include "common/moe_hybrid_routing_stats.h" #include "common/moe_hybrid_stream.h" #include "common/moe_router_graph.h" #include "internal.h" @@ -450,7 +451,8 @@ bool streamed_kimi_k3_step( std::vector & logits, MoeHybridStreamEngine & stream_engine, MoeStreamDualOwnerExecutor * dual_stream_executor, - const MoeStreamDualOwnerPolicy * stream_owner_policy) { + const MoeStreamDualOwnerPolicy * stream_owner_policy, + MoeHybridRoutingStats * routing_stats) { std::vector hidden(static_cast(w.n_embd)); { @@ -639,6 +641,14 @@ bool streamed_kimi_k3_step( route_batch.inputs = routed_input_host.data(); route_batch.selected_ids = selected.data(); route_batch.selected_weights = route_weights.data(); + if (routing_stats && !routing_stats->observe( + route_batch.layer, selected.data(), + static_cast(selected.size()))) { + set_last_error( + "Kimi-K3 routed layer " + std::to_string(il) + + ": failed to record native route statistics"); + return false; + } std::vector routed_output; std::string stream_error; MoeStreamDualOwnerStats owner_stats; @@ -822,7 +832,8 @@ bool kimi_k3_step(ggml_backend_t backend, std::vector & logits, MoeHybridStreamEngine * stream_engine, MoeStreamDualOwnerExecutor * dual_stream_executor, - const MoeStreamDualOwnerPolicy * stream_owner_policy) { + const MoeStreamDualOwnerPolicy * stream_owner_policy, + MoeHybridRoutingStats * routing_stats) { if (!backend || !w.ctx || !cache.ctx || position < 0 || position >= cache.max_ctx || position != cache.cur_pos || token < 0 || token >= w.n_vocab) { @@ -845,7 +856,7 @@ bool kimi_k3_step(ggml_backend_t backend, return streamed_kimi_k3_step( backend, w, cache, token, position, logits, *stream_engine, dual_stream_executor, - stream_owner_policy); + stream_owner_policy, routing_stats); } ggml_init_params params{}; diff --git a/server/src/kimi_k3/kimi_k3_internal.h b/server/src/kimi_k3/kimi_k3_internal.h index 1f1f9c48f..42d9f6cf7 100644 --- a/server/src/kimi_k3/kimi_k3_internal.h +++ b/server/src/kimi_k3/kimi_k3_internal.h @@ -23,6 +23,8 @@ namespace dflash::common { +struct MoeHybridRoutingStats; + class MoeHybridStreamEngine; class MoeStreamDualOwnerExecutor; struct MoeStreamDualOwnerPolicy; @@ -176,6 +178,7 @@ bool kimi_k3_step(ggml_backend_t backend, std::vector & logits, MoeHybridStreamEngine * stream_engine = nullptr, MoeStreamDualOwnerExecutor * dual_stream_executor = nullptr, - const MoeStreamDualOwnerPolicy * stream_owner_policy = nullptr); + const MoeStreamDualOwnerPolicy * stream_owner_policy = nullptr, + MoeHybridRoutingStats * routing_stats = nullptr); } // namespace dflash::common diff --git a/server/test/test_moe_expert_package.cpp b/server/test/test_moe_expert_package.cpp new file mode 100644 index 000000000..6e64a6add --- /dev/null +++ b/server/test/test_moe_expert_package.cpp @@ -0,0 +1,263 @@ +#include "CppUnitTestFramework.hpp" +#include "common/moe_expert_package.h" + +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#include +#include +#else +#include +#include +#include +#endif + +using namespace dflash::common; + +#define PACKAGE_REQUIRE(cond) do { \ + if (!(cond)) throw std::runtime_error(std::string(__FILE__) + ":" + \ + std::to_string(__LINE__) + ": " + #cond); \ +} while (0) + +namespace { + +struct MoeExpertPackageFixture {}; + +bool aligned_allocate(void ** ptr, size_t bytes, void *) { +#if defined(_WIN32) + *ptr = _aligned_malloc(bytes, 4096); + return *ptr != nullptr; +#else + return ::posix_memalign(ptr, 4096, bytes) == 0; +#endif +} + +void aligned_free(void * ptr, void *) { +#if defined(_WIN32) + _aligned_free(ptr); +#else + std::free(ptr); +#endif +} + +uint8_t pattern(int layer, int component, int expert, size_t byte) { + return (uint8_t) ((layer * 67 + component * 29 + expert * 13 + byte * 3) & 0xff); +} + +void fill_component(std::vector & source, + const ExpertFileRegion & region, + size_t expert_bytes, int layer, int component, + uint32_t experts) { + for (uint32_t expert = 0; expert < experts; ++expert) { + for (size_t byte = 0; byte < expert_bytes; ++byte) { + source[region.offset + (size_t) expert * expert_bytes + byte] = + pattern(layer, component, (int) expert, byte); + } + } +} + +struct PackageModel { + std::vector shard0 = std::vector(64 * 1024, 0xa5); + std::vector shard1 = std::vector(64 * 1024, 0x5a); + std::vector layers{2}; + std::vector experts{3, 2}; + + PackageModel() { + LayerExpertRegions & ordinary = layers[0]; + ordinary.expert_bytes_gate = 23; + ordinary.expert_bytes_up = 19; + ordinary.expert_bytes_down = 31; + ordinary.gate_exps = {101, ordinary.expert_bytes_gate * experts[0], 0}; + ordinary.up_exps = {211, ordinary.expert_bytes_up * experts[0], 1}; + ordinary.down_exps = {401, ordinary.expert_bytes_down * experts[0], 1}; + fill_component(shard0, ordinary.gate_exps, + ordinary.expert_bytes_gate, 0, 0, experts[0]); + fill_component(shard1, ordinary.up_exps, + ordinary.expert_bytes_up, 0, 1, experts[0]); + fill_component(shard1, ordinary.down_exps, + ordinary.expert_bytes_down, 0, 2, experts[0]); + + LayerExpertRegions & fused = layers[1]; + fused.fused_gate_up = true; + fused.expert_bytes_gate_up = 47; + fused.expert_bytes_down = 37; + fused.gate_up_exps = {1001, fused.expert_bytes_gate_up * experts[1], 0}; + fused.down_exps = {1401, fused.expert_bytes_down * experts[1], 1}; + fill_component(shard0, fused.gate_up_exps, + fused.expert_bytes_gate_up, 1, 0, experts[1]); + fill_component(shard1, fused.down_exps, + fused.expert_bytes_down, 1, 1, experts[1]); + } + + std::vector sources() const { + return { + {shard0.data(), shard0.size(), -1}, + {shard1.data(), shard1.size(), -1}, + }; + } +}; + +int make_temp_file(char * path) { +#if defined(_WIN32) + (void) path; + char directory[MAX_PATH]{}; + char generated[MAX_PATH]{}; + if (::GetTempPathA(MAX_PATH, directory) == 0 || + ::GetTempFileNameA(directory, "lbm", 0, generated) == 0) { + return -1; + } + return ::_open(generated, + _O_CREAT | _O_TRUNC | _O_RDWR | _O_BINARY | _O_TEMPORARY, + _S_IREAD | _S_IWRITE); +#else + const int fd = ::mkstemp(path); + if (fd >= 0) ::unlink(path); + return fd; +#endif +} + +void close_file(int fd) { +#if defined(_WIN32) + ::_close(fd); +#else + ::close(fd); +#endif +} + +size_t file_size(int fd) { +#if defined(_WIN32) + struct _stat64 st {}; + return ::_fstat64(fd, &st) == 0 ? (size_t) st.st_size : 0; +#else + struct stat st {}; + return ::fstat(fd, &st) == 0 ? (size_t) st.st_size : 0; +#endif +} + +void verify_component(const MoeNvmeLease & lease, + MoeExpertComponentKind kind, + int layer, int component, int expert, size_t bytes) { + const MoeExpertComponentLayout * layout = lease.layout().component(kind); + PACKAGE_REQUIRE(layout != nullptr); + const MoeExpertIoSpan & span = lease.layout().spans[0]; + PACKAGE_REQUIRE(layout->device_offset >= span.device_offset); + const size_t delta = layout->device_offset - span.device_offset; + const uint8_t * data = lease.data() + span.buffer_offset + delta; + PACKAGE_REQUIRE(data[0] == pattern(layer, component, expert, 0)); + PACKAGE_REQUIRE(data[bytes / 2] == pattern(layer, component, expert, bytes / 2)); + PACKAGE_REQUIRE(data[bytes - 1] == pattern(layer, component, expert, bytes - 1)); +} + +} // namespace + +TEST_CASE(MoeExpertPackageFixture, plans_aligned_one_read_records) { + PackageModel model; + MoeExpertPackageManifest manifest; + MoeExpertPackageOptions options; + options.sync_on_finish = false; + std::string err; + PACKAGE_REQUIRE(plan_moe_expert_package( + model.sources(), model.layers, model.experts, + options, manifest, &err)); + PACKAGE_REQUIRE(manifest.version == 1); + PACKAGE_REQUIRE(manifest.layer_regions.size() == 2); + PACKAGE_REQUIRE(manifest.source_layout_hash != 0); + for (const LayerExpertRegions & layer : manifest.layer_regions) { + PACKAGE_REQUIRE(layer.expert_major.enabled); + PACKAGE_REQUIRE(layer.expert_major.expert_stride % 4096 == 0); + PACKAGE_REQUIRE(layer.expert_major.experts.offset % 4096 == 0); + PACKAGE_REQUIRE(layer.expert_major.experts.source_index == 0); + } +} + +TEST_CASE(MoeExpertPackageFixture, roundtrip_preserves_split_source_bytes) { + PackageModel model; + char path[] = "/tmp/moe_expert_package_XXXXXX"; + const int fd = make_temp_file(path); + PACKAGE_REQUIRE(fd >= 0); + + MoeExpertPackageOptions options; + options.sync_on_finish = false; + MoeExpertPackageManifest written; + std::string err; + PACKAGE_REQUIRE(write_moe_expert_package( + fd, model.sources(), model.layers, model.experts, + options, &written, &err)); + PACKAGE_REQUIRE(file_size(fd) == written.file_bytes); + + MoeNvmeSource package{nullptr, file_size(fd), fd}; + MoeExpertPackageManifest loaded; + PACKAGE_REQUIRE(read_moe_expert_package(package, loaded, &err)); + PACKAGE_REQUIRE(loaded.source_layout_hash == + moe_expert_source_layout_hash( + model.sources(), model.layers, model.experts)); + PACKAGE_REQUIRE(loaded.expert_counts == model.experts); + + MoeNvmeConfig config; + config.backend = MoeNvmeBackend::ThreadPool; + config.direct_io = MoeNvmeDirectMode::Disabled; + config.host_slots = 2; + config.io_threads = 1; + MoeNvmeScheduler scheduler; + PACKAGE_REQUIRE(scheduler.init( + config, loaded.max_record_bytes, + aligned_allocate, aligned_free, nullptr, &err)); + PACKAGE_REQUIRE(scheduler.bind_source( + package, loaded.layer_regions, &err)); + + MoeNvmeLease ordinary; + PACKAGE_REQUIRE(scheduler.acquire(0, 2, ordinary, &err)); + PACKAGE_REQUIRE(ordinary.layout().span_count == 1); + verify_component(ordinary, MoeExpertComponentKind::Gate, + 0, 0, 2, model.layers[0].expert_bytes_gate); + verify_component(ordinary, MoeExpertComponentKind::Up, + 0, 1, 2, model.layers[0].expert_bytes_up); + verify_component(ordinary, MoeExpertComponentKind::Down, + 0, 2, 2, model.layers[0].expert_bytes_down); + ordinary.reset(); + + MoeNvmeLease fused; + PACKAGE_REQUIRE(scheduler.acquire(1, 1, fused, &err)); + PACKAGE_REQUIRE(fused.layout().span_count == 1); + verify_component(fused, MoeExpertComponentKind::FusedGateUp, + 1, 0, 1, model.layers[1].expert_bytes_gate_up); + verify_component(fused, MoeExpertComponentKind::Down, + 1, 1, 1, model.layers[1].expert_bytes_down); + fused.reset(); + PACKAGE_REQUIRE(scheduler.stats().read_ops == 2); + PACKAGE_REQUIRE(scheduler.stats().errors == 0); + scheduler.destroy(); + close_file(fd); +} + +TEST_CASE(MoeExpertPackageFixture, rejects_incomplete_and_mismatched_sources) { + PackageModel model; + const uint64_t original_hash = moe_expert_source_layout_hash( + model.sources(), model.layers, model.experts); + PACKAGE_REQUIRE(original_hash != 0); + model.shard0[model.layers[0].gate_exps.offset] ^= 0x1; + PACKAGE_REQUIRE(moe_expert_source_layout_hash( + model.sources(), model.layers, model.experts) != original_hash); + + std::vector sources = model.sources(); + sources[1].mmap_size = 256; + MoeExpertPackageManifest manifest; + MoeExpertPackageOptions options; + options.sync_on_finish = false; + std::string err; + PACKAGE_REQUIRE(!plan_moe_expert_package( + sources, model.layers, model.experts, + options, manifest, &err)); + PACKAGE_REQUIRE(!err.empty()); + + std::vector incomplete(4096, 0); + PACKAGE_REQUIRE(!read_moe_expert_package( + {incomplete.data(), incomplete.size(), -1}, manifest, &err)); + PACKAGE_REQUIRE(err.find("magic") != std::string::npos); +} diff --git a/server/test/test_moe_stream_compute.cpp b/server/test/test_moe_stream_compute.cpp index bbf2ab7cb..7b3a2ad48 100644 --- a/server/test/test_moe_stream_compute.cpp +++ b/server/test/test_moe_stream_compute.cpp @@ -436,6 +436,76 @@ void run_mxfp4_padding_case(ggml_backend_t backend) { engine.destroy(); } +void run_pinned_cache_case(ggml_backend_t backend) { + std::vector gate; + std::vector up; + std::vector down; + fill_weights(gate, up, down); + ModelBytes model = make_model_bytes(false, gate, up, down); + TempFile file(model.file); + + MoeHybridStorage storage; + storage.mmap_size = model.file.size(); + storage.mmap_fd = ::dup(file.fd); + STREAM_REQUIRE(storage.mmap_fd >= 0); + // A second logical layer gives the eviction test more unique cache keys + // without making the synthetic weight file larger. + storage.layer_regions = {model.regions, model.regions}; + + MoeStreamConfig config; + config.device_slots = 3; + config.device_cache_bytes = 0; + config.graph_cache_entries = 0; + config.nvme.backend = MoeNvmeBackend::ThreadPool; + config.nvme.direct_io = MoeNvmeDirectMode::Disabled; + config.nvme.host_slots = 6; + config.nvme.io_threads = 2; + + MoeHybridStreamEngine engine; + std::string error; + STREAM_REQUIRE(engine.init( + backend, model.slot_bytes, storage, config, &error)); + STREAM_REQUIRE(engine.device_slot_count() == 3); + + MoeStreamExpertSpec spec; + spec.input_dim = kInput; + spec.intermediate_dim = kFf; + spec.output_dim = kOutput; + spec.gate_type = GGML_TYPE_F32; + spec.up_type = GGML_TYPE_F32; + spec.down_type = GGML_TYPE_F32; + spec.gated_activation = MoeGatedActivation::Situ; + + const std::vector layer_specs = {spec, spec}; + const std::vector warm = { + {0, 0, 100, model.slot_bytes}, + }; + MoeStreamCacheWarmStats warm_stats; + STREAM_REQUIRE(engine.warm_and_pin_device_cache( + layer_specs, warm, 2, &warm_stats, &error)); + STREAM_REQUIRE(warm_stats.admitted == 1); + STREAM_REQUIRE(engine.pinned_expert_count() == 1); + + auto touch = [&](int layer, int expert) { + int slot = -1; + STREAM_REQUIRE(engine.stage_expert_cached_async( + layer, expert, &slot, &error)); + STREAM_REQUIRE(engine.activate_device_slot(slot, &error)); + engine.release_device_slot(slot); + }; + // Only two slots are evictable. Four distinct cold keys force churn. + touch(0, 1); + touch(0, 2); + touch(1, 0); + touch(1, 1); + + const uint64_t requests_before_hot_reuse = engine.io_stats().requests; + touch(0, 0); + STREAM_REQUIRE(engine.io_stats().requests == requests_before_hot_reuse); + STREAM_REQUIRE(engine.pinned_expert_count() == 1); + engine.destroy(); +} + } // namespace TEST_CASE(MoeStreamComputeFixture, persistent_graph_matches_cpu_and_padded_mxfp4) { @@ -455,5 +525,6 @@ TEST_CASE(MoeStreamComputeFixture, persistent_graph_matches_cpu_and_padded_mxfp4 run_layout_case(backend, false); run_layout_case(backend, true); run_mxfp4_padding_case(backend); + run_pinned_cache_case(backend); ggml_backend_free(backend); } diff --git a/server/test/test_moe_stream_owner_partition.cpp b/server/test/test_moe_stream_owner_partition.cpp index 52e1ca1fc..2862f3623 100644 --- a/server/test/test_moe_stream_owner_partition.cpp +++ b/server/test/test_moe_stream_owner_partition.cpp @@ -1,5 +1,6 @@ #include "CppUnitTestFramework.hpp" #include "../src/common/moe_hybrid_stream.h" +#include "../src/common/moe_stream_cache_policy.h" #include #include @@ -136,3 +137,53 @@ TEST_CASE(MoeStreamOwnerPartitionFixture, REQUIRE(!partition_moe_stream_routes( batch, policy, primary, secondary, nullptr, &error)); } + +TEST_CASE(MoeStreamOwnerPartitionFixture, + warm_cache_plan_is_value_ranked_and_owner_disjoint) { + MoeHybridRoutingStats routing; + REQUIRE(routing.init(2, 4, 2)); + routing.counts = { + 100, 40, 20, 10, + 80, 70, 60, 50, + }; + routing.layer_totals = {170, 260}; + + MoeHybridPlacement placement; + placement.n_layer = 2; + placement.n_expert = 4; + placement.n_expert_used = 2; + placement.total_hot = 4; + placement.hot_counts = {2, 2}; + placement.hot_expert_ids = {{0, 2}, {1, 3}}; + MoeStreamDualOwnerPolicy owner_policy; + owner_policy.primary_placement = &placement; + + MoeStreamCachePlanConfig config; + config.max_entries = 3; + config.max_bytes = 400; + config.owner = MoeStreamCacheOwner::Primary; + std::vector primary; + std::string error; + REQUIRE(build_moe_stream_cache_plan( + routing, {100, 200}, config, &owner_policy, + primary, &error)); + REQUIRE(primary.size() == 3); + REQUIRE(primary[0].layer == 0 && primary[0].expert == 0); + REQUIRE(primary[1].layer == 1 && primary[1].expert == 1); + REQUIRE(primary[2].layer == 0 && primary[2].expert == 2); + + config.owner = MoeStreamCacheOwner::Secondary; + std::vector secondary; + REQUIRE(build_moe_stream_cache_plan( + routing, {100, 200}, config, &owner_policy, + secondary, &error)); + REQUIRE(!secondary.empty()); + for (const MoeStreamCacheWarmEntry & entry : primary) { + REQUIRE(moe_stream_primary_owns_expert( + owner_policy, entry.layer, entry.expert)); + } + for (const MoeStreamCacheWarmEntry & entry : secondary) { + REQUIRE(!moe_stream_primary_owns_expert( + owner_policy, entry.layer, entry.expert)); + } +} From 3450da3fbcc18763209670a875707b1d3f96f11c Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 2 Aug 2026 17:24:38 +0200 Subject: [PATCH 15/20] perf(moe): fuse resident decode expert graphs --- server/docs/KIMI_K3_HETERO.md | 29 +- server/docs/MOE_NVME_STREAMING.md | 33 +- server/src/common/moe_hybrid_stream.cpp | 577 ++++++++++++++++++++---- server/src/common/moe_hybrid_stream.h | 7 + server/test/bench_kimi_k3_hetero.cpp | 7 +- server/test/test_moe_stream_compute.cpp | 151 ++++++- 6 files changed, 700 insertions(+), 104 deletions(-) diff --git a/server/docs/KIMI_K3_HETERO.md b/server/docs/KIMI_K3_HETERO.md index e69a9097f..d0f196a5c 100644 --- a/server/docs/KIMI_K3_HETERO.md +++ b/server/docs/KIMI_K3_HETERO.md @@ -42,6 +42,11 @@ small architecture adapter around it: 896 experts, top-16, 92 MoE layers, `3584 -> 3072 -> 3584`, and SiTU (`beta=4`, `linear_beta=25`). It overlaps SSD/H2D for expert N+1 with compute for expert N. +- All-resident single-token route sets use one persistent GPU fork/join graph + per MoE layer. The graph evaluates every selected expert, applies dynamic + native-router weights, and reduces on-device, replacing up to 16 graph + submissions, synchronizations, and output copies. Any cache miss retains the + original overlap pipeline automatically. Seven scheduler tests pass on the AMD Lucebox, including expert-major one-read records, mmap, and real-file multi-shard reads. A separate numerical test @@ -51,8 +56,8 @@ source index zero and need no model-specific change. The SSD text path is implemented, but two qualification boundaries remain: -- The backend is correctness-first and token-sequential. Its per-layer graph - boundaries are not yet fused/captured for full-model speed. +- The backend is correctness-first and token-sequential. Routed expert + fork/join is fused on cache hits, but full-model layer boundaries are not. - The heterogeneous routed branches overlap, but their partial outputs still cross a host-visible, activation-sized boundary at every routed layer. It does not yet use PR-505's device-resident peer join. @@ -129,6 +134,26 @@ The repeated-route result is deliberately a best case, not a prediction. Kimi K3 was designed for balanced expert use. With unrelated balanced routes, a cache only helps in proportion to its share of the 495 GiB routed pool. +## Warm routed-core optimization, 2026-08-02 + +Lucebox3 supplied independent Strix Halo and RTX 3090 qualification. The +benchmark used the complete Kimi geometry above, 50 tokens with identical +routes, and a 9.763 GiB device cache. The first token was cold and moved the +exact 8.843994 GiB; the remaining tokens were device-cache hits. + +| Backend | Per-expert graphs | Fused resident decode | Gain | +|---|---:|---:|---:| +| Strix Halo ROCm/gfx1151 (three-run mean) | 5.770 token/s | **6.913 token/s** | **+19.8%** | +| RTX 3090 CUDA/sm_86 (three-run mean) | 6.327 token/s | **7.364 token/s** | **+16.4%** | + +On Strix, graph submissions fell from 73,600 to 5,980: 1,472 submissions for +the first cold token plus one fused submission for each of the remaining +4,508 layer calls. A separate 0%-hit, 128 MiB-cache run selected the original +cold pipeline and issued zero fused launches. Numerical tests match a CPU +oracle on both backends and also change expert addresses and route weights +between cached graph launches. These numbers isolate the routed core and do +not claim complete-model token speed. + ## Practical Lucebox placement The machine has about 125.08 GiB of system/UMA memory plus 31.86 GiB on the diff --git a/server/docs/MOE_NVME_STREAMING.md b/server/docs/MOE_NVME_STREAMING.md index 4d732597e..c2dfebb44 100644 --- a/server/docs/MOE_NVME_STREAMING.md +++ b/server/docs/MOE_NVME_STREAMING.md @@ -87,6 +87,11 @@ weight masks reconstruct the original route batch exactly. - Persistent compute graphs remove graph construction and activation-buffer allocation from the steady-state expert loop. A bounded LRU supports models whose layers use more than one expert shape or quantization format. +- When every route selected for a decode token is already device-resident, one + persistent fork/join graph evaluates all selected experts and performs the + weighted reduction on the GPU. A cache miss automatically retains the + transfer/compute pipeline, so this optimization cannot remove cold-path + overlap. The native model router remains authoritative. Prediction may only issue a bounded prefetch; a wrong prediction cannot change model output. @@ -169,9 +174,10 @@ budget reserves the larger of 2 GiB or 5% of device memory and never exceeds the complete routed pool. `DFLASH_MOE_NVME_DEVICE_CACHE_MB` remains an explicit override, also capped by the routed pool. -Kimi's native router remains authoritative. The current text backend is -correctness-first and sequential; captured per-layer graphs and the vision -tower are separate optimizations. +Kimi's native router remains authoritative. Routed decode uses one GPU graph +per layer when every selected expert is resident; otherwise it uses the +pipelined per-expert graph. The complete text backend remains layer-sequential, +and the vision tower is a separate optimization. ### Expert-major package and profile-guided cache @@ -259,6 +265,7 @@ not improve the qualified P310 drive and consumes extra pinned/system memory. | `DFLASH_MOE_NVME_DEVICE_SLOTS` | `2` | Minimum rotating GPU expert buffers | | `DFLASH_MOE_NVME_DEVICE_CACHE_MB` | automatic/`0` | Override adaptive device expert-cache memory; `0` leaves only pipeline slots | | `DFLASH_MOE_NVME_GRAPH_CACHE` | `8` | Persistent expert-graph variants retained per stream engine; `0` is a diagnostic no-cache mode | +| `DFLASH_MOE_NVME_FUSED_DECODE` | `1` | Fuse an all-resident single-token route set and its weighted reduction into one GPU graph; `0` is an A/B fallback | | `DFLASH_MOE_NVME_REFERENCE_EVAL` | unset | Diagnostic only: `1` restores the allocation-heavy reference evaluator for numerical/performance A/B | | `DFLASH_MOE_EXPERT_PACKAGE` | unset | Optional validated expert-major package used in place of tensor-major GGUF expert regions | | `DFLASH_MOE_EXPERT_PACKAGE_BUILD` | unset | One-time Kimi package creation: `1` builds when absent; `force` rebuilds | @@ -271,7 +278,8 @@ not improve the qualified P310 drive and consumes extra pinned/system memory. Shutdown telemetry reports logical and physical bytes, measured read service rate, cache hits, demand wait/timeouts, de-duplication, dropped speculation, -errors, and persistent-graph builds/hits/evictions. The scheduler rejects +errors, persistent-graph builds/hits/evictions, and fused decode +launches/experts. The scheduler rejects truncated shards at bind time and accepts a valid short direct-I/O completion only when it covers the complete logical payload at an unaligned file tail. `test_moe_stream_compute` generates @@ -283,6 +291,23 @@ targets `test_moe_nvme_scheduler`, `test_moe_expert_package`, `bench_moe_nvme_pipeline` test scheduling, raw storage, and the complete SSD-to-GPU path. Benchmarks are read-only. +### Warm decode qualification (2026-08-02) + +The model-neutral Kimi geometry benchmark was run on Lucebox3 with 896 IQ1_S +experts, top-16 routing, 92 MoE layers, and a 9.763 GiB device cache. The first +token loaded 8.843994 GiB; the following 49 tokens reused identical routes, for +a 98% aggregate cache-hit rate. Three alternating A/B runs on Strix Halo +measured 5.770 routed-core token/s with per-expert submission and 6.913 token/s +with fused decode: **+19.8%**. Each fused run reduced 73,600 expert graph +submissions to 1,472 cold submissions plus 4,508 fused layer submissions. + +The same numerical test passed on ROCm/gfx1151 and CUDA/sm_86, including graph +reuse with changed expert pointers and route weights. RTX 3090 throughput on +the same three-run A/B improved from 6.327 to 7.364 token/s (**+16.4%**). With a +128 MiB cache and unrelated routes, the resident guard selected the existing +cold pipeline (zero fused launches), preserving its storage throughput. These +are routed-core microbenchmarks rather than complete Kimi generation rates. + The external `smoke_kimi_k3_forward` target accepts `[stream_experts=0|1] [expert_gpu=-1]`. This is an A/B and placement oracle for small Kimi fixtures; production Kimi uses streaming and `-1` resolves the diff --git a/server/src/common/moe_hybrid_stream.cpp b/server/src/common/moe_hybrid_stream.cpp index 9ae234560..3d3af5bd3 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -144,6 +144,8 @@ MoeStreamConfig MoeStreamConfig::from_env() { "DFLASH_MOE_NVME_DEVICE_SLOTS", config.device_slots, 2, 8); config.graph_cache_entries = env_bounded_int( "DFLASH_MOE_NVME_GRAPH_CACHE", config.graph_cache_entries, 0, 64); + config.fused_decode = env_bounded_int( + "DFLASH_MOE_NVME_FUSED_DECODE", config.fused_decode ? 1 : 0, 0, 1) != 0; config.device_cache_bytes = env_mib( "DFLASH_MOE_NVME_DEVICE_CACHE_MB", config.device_cache_bytes); config.prefill_threshold = env_bounded_int( @@ -289,6 +291,111 @@ bool validate_moe_stream_expert_layout( namespace { +struct StreamExpertDeviceBinding { + const void * gate = nullptr; + const void * up = nullptr; + const void * down = nullptr; + size_t gate_alloc_bytes = 0; + size_t up_alloc_bytes = 0; + size_t down_alloc_bytes = 0; +}; + +bool bind_stream_tensor(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + const void * data, + size_t available_bytes, + const char * label, + std::string * err) { + if (!buffer || !tensor || !data) { + if (err) *err = std::string("invalid streamed ") + + label + " tensor binding"; + return false; + } + const size_t required_bytes = + ggml_backend_buffer_get_alloc_size(buffer, tensor); + if (available_bytes < required_bytes) { + if (err) *err = std::string("streamed ") + label + + " device allocation is smaller than the backend's padded " + "tensor requirement"; + return false; + } + const size_t alignment = ggml_backend_buffer_get_alignment(buffer); + if (alignment != 0 && (uintptr_t) data % alignment != 0) { + if (err) *err = std::string("streamed ") + label + + " tensor is not aligned for the compute backend"; + return false; + } + if (ggml_backend_tensor_alloc( + buffer, tensor, const_cast(data)) != GGML_STATUS_SUCCESS) { + if (err) *err = std::string("failed to bind streamed ") + label + + " tensor to expert cache"; + return false; + } + return true; +} + +ggml_tensor * scale_stream_tensor(ggml_context * ctx, + ggml_tensor * value, + float scale) { + return scale == 1.0f ? value : ggml_scale(ctx, value, scale); +} + +ggml_tensor * apply_stream_gated_activation( + ggml_context * ctx, + const MoeStreamExpertSpec & spec, + ggml_tensor * gate, + ggml_tensor * up) { + if (spec.gated_activation == MoeGatedActivation::Situ) { + ggml_tensor * nonlinear = ggml_scale(ctx, gate, 1.0f / spec.situ_beta); + nonlinear = ggml_tanh(ctx, nonlinear); + nonlinear = ggml_scale(ctx, nonlinear, spec.situ_beta); + nonlinear = ggml_mul(ctx, nonlinear, ggml_sigmoid(ctx, gate)); + ggml_tensor * linear = ggml_scale( + ctx, up, 1.0f / spec.situ_linear_beta); + linear = ggml_tanh(ctx, linear); + linear = ggml_scale(ctx, linear, spec.situ_linear_beta); + return ggml_mul(ctx, nonlinear, linear); + } + if (spec.swiglu_clamp > 0.0f) { + return ggml_swiglu_ds4_split(ctx, gate, up, spec.swiglu_clamp); + } + return ggml_swiglu_split(ctx, gate, up); +} + +ggml_tensor * build_stream_expert_branch( + ggml_context * ctx, + const MoeStreamExpertSpec & spec, + int batch, + ggml_tensor * input, + ggml_tensor * gate, + ggml_tensor * up, + ggml_tensor * down, + ggml_tensor * gate_up) { + ggml_tensor * activated = nullptr; + if (gate_up) { + ggml_tensor * combined = scale_stream_tensor( + ctx, ggml_mul_mat(ctx, gate_up, input), spec.gate_up_scale); + ggml_tensor * gate_part = ggml_view_2d( + ctx, combined, spec.intermediate_dim, batch, + combined->nb[1], 0); + ggml_tensor * up_part = ggml_view_2d( + ctx, combined, spec.intermediate_dim, batch, + combined->nb[1], + (size_t) spec.intermediate_dim * sizeof(float)); + activated = apply_stream_gated_activation( + ctx, spec, ggml_cont(ctx, gate_part), ggml_cont(ctx, up_part)); + } else { + ggml_tensor * gate_value = scale_stream_tensor( + ctx, ggml_mul_mat(ctx, gate, input), spec.gate_scale); + ggml_tensor * up_value = scale_stream_tensor( + ctx, ggml_mul_mat(ctx, up, input), spec.up_scale); + activated = apply_stream_gated_activation( + ctx, spec, gate_value, up_value); + } + return scale_stream_tensor( + ctx, ggml_mul_mat(ctx, down, activated), spec.down_scale); +} + class PersistentStreamExpertGraph { public: ~PersistentStreamExpertGraph() { destroy(); } @@ -350,99 +457,24 @@ class PersistentStreamExpertGraph { ggml_set_input(down_); } - auto bind_external = [&](ggml_tensor * tensor, - const void * data, - size_t available_bytes, - const char * label) -> bool { - if (!tensor || !data) { - if (err) *err = std::string("invalid streamed ") + - label + " tensor binding"; - return false; - } - const size_t required_bytes = - ggml_backend_buffer_get_alloc_size(expert_buffer, tensor); - if (available_bytes < required_bytes) { - if (err) *err = std::string("streamed ") + label + - " device allocation is smaller than the backend's " - "padded tensor requirement"; - return false; - } - const size_t alignment = - ggml_backend_buffer_get_alignment(expert_buffer); - if (alignment != 0 && - (uintptr_t) data % alignment != 0) { - if (err) *err = std::string("streamed ") + label + - " tensor is not aligned for the compute backend"; - return false; - } - if (ggml_backend_tensor_alloc( - expert_buffer, tensor, const_cast(data)) != - GGML_STATUS_SUCCESS) { - if (err) *err = std::string("failed to bind streamed ") + - label + " tensor to expert cache"; - return false; - } - return true; - }; if (spec.fused_gate_up) { - if (!bind_external(gate_up_, gate_data, gate_alloc_bytes, - "gate_up") || - !bind_external(down_, down_data, down_alloc_bytes, "down")) { + if (!bind_stream_tensor(expert_buffer, gate_up_, gate_data, + gate_alloc_bytes, "gate_up", err) || + !bind_stream_tensor(expert_buffer, down_, down_data, + down_alloc_bytes, "down", err)) { return false; } - } else if (!bind_external(gate_, gate_data, gate_alloc_bytes, "gate") || - !bind_external(up_, up_data, up_alloc_bytes, "up") || - !bind_external(down_, down_data, down_alloc_bytes, "down")) { + } else if (!bind_stream_tensor(expert_buffer, gate_, gate_data, + gate_alloc_bytes, "gate", err) || + !bind_stream_tensor(expert_buffer, up_, up_data, + up_alloc_bytes, "up", err) || + !bind_stream_tensor(expert_buffer, down_, down_data, + down_alloc_bytes, "down", err)) { return false; } - auto scale_if_needed = [&](ggml_tensor * value, float scale) { - return scale == 1.0f ? value : ggml_scale(ctx_, value, scale); - }; - auto gated_activation = [&](ggml_tensor * gate, - ggml_tensor * up) -> ggml_tensor * { - if (spec.gated_activation == MoeGatedActivation::Situ) { - ggml_tensor * nonlinear = ggml_scale( - ctx_, gate, 1.0f / spec.situ_beta); - nonlinear = ggml_tanh(ctx_, nonlinear); - nonlinear = ggml_scale(ctx_, nonlinear, spec.situ_beta); - nonlinear = ggml_mul( - ctx_, nonlinear, ggml_sigmoid(ctx_, gate)); - ggml_tensor * linear = ggml_scale( - ctx_, up, 1.0f / spec.situ_linear_beta); - linear = ggml_tanh(ctx_, linear); - linear = ggml_scale(ctx_, linear, spec.situ_linear_beta); - return ggml_mul(ctx_, nonlinear, linear); - } - if (spec.swiglu_clamp > 0.0f) { - return ggml_swiglu_ds4_split( - ctx_, gate, up, spec.swiglu_clamp); - } - return ggml_swiglu_split(ctx_, gate, up); - }; - - ggml_tensor * activated = nullptr; - if (gate_up_) { - ggml_tensor * combined = scale_if_needed( - ggml_mul_mat(ctx_, gate_up_, input_), spec.gate_up_scale); - ggml_tensor * gate_part = ggml_view_2d( - ctx_, combined, spec.intermediate_dim, batch, - combined->nb[1], 0); - ggml_tensor * up_part = ggml_view_2d( - ctx_, combined, spec.intermediate_dim, batch, - combined->nb[1], - (size_t) spec.intermediate_dim * sizeof(float)); - activated = gated_activation( - ggml_cont(ctx_, gate_part), ggml_cont(ctx_, up_part)); - } else { - ggml_tensor * gate_value = scale_if_needed( - ggml_mul_mat(ctx_, gate_, input_), spec.gate_scale); - ggml_tensor * up_value = scale_if_needed( - ggml_mul_mat(ctx_, up_, input_), spec.up_scale); - activated = gated_activation(gate_value, up_value); - } - output_ = scale_if_needed( - ggml_mul_mat(ctx_, down_, activated), spec.down_scale); + output_ = build_stream_expert_branch( + ctx_, spec, batch, input_, gate_, up_, down_, gate_up_); ggml_set_output(output_); graph_ = ggml_new_graph_custom(ctx_, 512, false); ggml_build_forward_expand(graph_, output_); @@ -535,6 +567,212 @@ class PersistentStreamExpertGraph { ggml_tensor * output_ = nullptr; }; +// Single-token decode selects several independent experts and then computes a +// weighted sum. Keeping those branches in one persistent graph removes the +// host boundary between experts: one graph submission, one synchronization, +// and one output copy per routed MoE layer. +class PersistentStreamMoEDecodeGraph { +public: + ~PersistentStreamMoEDecodeGraph() { destroy(); } + + bool matches(const MoeStreamExpertSpec & spec, int expert_count) const { + return expert_count_ == expert_count && same_stream_spec(spec_, spec); + } + + bool build(ggml_backend_t backend, + ggml_backend_buffer_t expert_buffer, + const MoeStreamExpertSpec & spec, + const std::vector & bindings, + std::string * err) { + destroy(); + if (!backend || !expert_buffer || bindings.size() < 2) { + if (err) *err = "invalid fused streamed-MoE graph arguments"; + return false; + } + backend_ = backend; + spec_ = spec; + expert_count_ = (int) bindings.size(); + + ggml_init_params params{}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ctx_ = ggml_init(params); + if (!ctx_) { + if (err) *err = "ggml_init failed for fused streamed-MoE decode"; + return false; + } + + input_ = ggml_new_tensor_2d(ctx_, GGML_TYPE_F32, spec.input_dim, 1); + route_weights_ = ggml_new_tensor_1d( + ctx_, GGML_TYPE_F32, expert_count_); + ggml_set_input(input_); + ggml_set_input(route_weights_); + + tensors_.resize(bindings.size()); + expert_outputs_ = ggml_new_tensor_3d( + ctx_, GGML_TYPE_F32, spec.output_dim, expert_count_, 1); + copy_nodes_.reserve(bindings.size()); + for (size_t i = 0; i < bindings.size(); ++i) { + ExpertTensors & tensors = tensors_[i]; + if (spec.fused_gate_up) { + tensors.gate_up = ggml_new_tensor_2d( + ctx_, spec.gate_up_type, spec.input_dim, + 2LL * spec.intermediate_dim); + tensors.down = ggml_new_tensor_2d( + ctx_, spec.down_type, spec.intermediate_dim, + spec.output_dim); + ggml_set_input(tensors.gate_up); + ggml_set_input(tensors.down); + if (!bind_stream_tensor( + expert_buffer, tensors.gate_up, bindings[i].gate, + bindings[i].gate_alloc_bytes, "gate_up", err) || + !bind_stream_tensor( + expert_buffer, tensors.down, bindings[i].down, + bindings[i].down_alloc_bytes, "down", err)) { + return false; + } + } else { + tensors.gate = ggml_new_tensor_2d( + ctx_, spec.gate_type, spec.input_dim, + spec.intermediate_dim); + tensors.up = ggml_new_tensor_2d( + ctx_, spec.up_type, spec.input_dim, + spec.intermediate_dim); + tensors.down = ggml_new_tensor_2d( + ctx_, spec.down_type, spec.intermediate_dim, + spec.output_dim); + ggml_set_input(tensors.gate); + ggml_set_input(tensors.up); + ggml_set_input(tensors.down); + if (!bind_stream_tensor( + expert_buffer, tensors.gate, bindings[i].gate, + bindings[i].gate_alloc_bytes, "gate", err) || + !bind_stream_tensor( + expert_buffer, tensors.up, bindings[i].up, + bindings[i].up_alloc_bytes, "up", err) || + !bind_stream_tensor( + expert_buffer, tensors.down, bindings[i].down, + bindings[i].down_alloc_bytes, "down", err)) { + return false; + } + } + + ggml_tensor * branch = build_stream_expert_branch( + ctx_, spec, 1, input_, tensors.gate, tensors.up, + tensors.down, tensors.gate_up); + ggml_tensor * destination = ggml_view_2d( + ctx_, expert_outputs_, spec.output_dim, 1, + expert_outputs_->nb[1], i * expert_outputs_->nb[1]); + ggml_tensor * copy = ggml_cpy(ctx_, branch, destination); + ggml_set_output(copy); + copy_nodes_.push_back(copy); + } + + output_ = ggml_laguna_moe_combine( + ctx_, expert_outputs_, route_weights_); + ggml_set_output(output_); + graph_ = ggml_new_graph_custom( + ctx_, std::max(512, bindings.size() * 256), false); + for (ggml_tensor * copy : copy_nodes_) { + ggml_build_forward_expand(graph_, copy); + } + ggml_build_forward_expand(graph_, output_); + alloc_ = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend_)); + if (!alloc_ || !ggml_gallocr_alloc_graph(alloc_, graph_)) { + if (err) *err = "fused streamed-MoE graph allocation failed"; + return false; + } + return true; + } + + bool launch(const std::vector & bindings, + const float * input, + const float * route_weights, + std::string * err) { + if (!valid() || !input || !route_weights || + bindings.size() != tensors_.size()) { + if (err) *err = "fused streamed-MoE graph is not ready"; + return false; + } + for (size_t i = 0; i < bindings.size(); ++i) { + ExpertTensors & tensors = tensors_[i]; + if (tensors.gate_up) { + tensors.gate_up->data = const_cast(bindings[i].gate); + } else { + tensors.gate->data = const_cast(bindings[i].gate); + tensors.up->data = const_cast(bindings[i].up); + } + tensors.down->data = const_cast(bindings[i].down); + } + ggml_backend_tensor_set( + input_, input, 0, (size_t) spec_.input_dim * sizeof(float)); + ggml_backend_tensor_set( + route_weights_, route_weights, 0, + bindings.size() * sizeof(float)); + if (ggml_backend_graph_compute_async(backend_, graph_) != + GGML_STATUS_SUCCESS) { + if (err) *err = "fused streamed-MoE graph launch failed"; + return false; + } + return true; + } + + bool finish(std::vector & output, std::string * err) { + if (!valid()) { + if (err) *err = "fused streamed-MoE graph is not ready"; + return false; + } + ggml_backend_synchronize(backend_); + output.resize((size_t) spec_.output_dim); + ggml_backend_tensor_get( + output_, output.data(), 0, + output.size() * sizeof(float)); + return true; + } + + void destroy() { + if (alloc_) ggml_gallocr_free(alloc_); + alloc_ = nullptr; + if (ctx_) ggml_free(ctx_); + ctx_ = nullptr; + graph_ = nullptr; + input_ = route_weights_ = expert_outputs_ = output_ = nullptr; + tensors_.clear(); + copy_nodes_.clear(); + backend_ = nullptr; + expert_count_ = 0; + } + + bool valid() const { + return backend_ && ctx_ && graph_ && alloc_ && input_ && + route_weights_ && expert_outputs_ && output_; + } + + uint64_t last_touch = 0; + +private: + struct ExpertTensors { + ggml_tensor * gate = nullptr; + ggml_tensor * up = nullptr; + ggml_tensor * down = nullptr; + ggml_tensor * gate_up = nullptr; + }; + + ggml_backend_t backend_ = nullptr; + MoeStreamExpertSpec spec_{}; + int expert_count_ = 0; + ggml_context * ctx_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t alloc_ = nullptr; + ggml_tensor * input_ = nullptr; + ggml_tensor * route_weights_ = nullptr; + ggml_tensor * expert_outputs_ = nullptr; + ggml_tensor * output_ = nullptr; + std::vector tensors_; + std::vector copy_nodes_; +}; + } // namespace struct MoeHybridStreamEngine::Runtime { @@ -597,6 +835,8 @@ struct MoeHybridStreamEngine::Runtime { size_t pinned_experts = 0; int active_slot = -1; std::vector> graph_cache; + std::vector> + fused_decode_graph_cache; uint64_t graph_clock = 0; MoeStreamComputeStats compute_stats{}; std::mutex compute_mutex; @@ -609,6 +849,7 @@ void release_device_cache(RuntimeT & runtime) { (void) cudaStreamSynchronize(runtime.transfer_stream); } runtime.graph_cache.clear(); + runtime.fused_decode_graph_cache.clear(); for (auto & slot : runtime.device_slots) { slot.host_lease.reset(); if (slot.ready) (void) cudaEventDestroy(slot.ready); @@ -938,6 +1179,7 @@ void MoeHybridStreamEngine::destroy() { (void) cudaStreamSynchronize(runtime_->transfer_stream); } runtime_->graph_cache.clear(); + runtime_->fused_decode_graph_cache.clear(); for (Runtime::DeviceSlot & slot : runtime_->device_slots) { slot.host_lease.reset(); if (slot.ready) (void) cudaEventDestroy(slot.ready); @@ -977,7 +1219,8 @@ void MoeHybridStreamEngine::destroy() { "timeouts=%llu errors=%llu " "device-cache=%.1f MiB slots=%zu pinned=%zu " "hits=%llu misses=%llu evictions=%llu " - "graphs=%llu graph-hits=%llu graph-evictions=%llu launches=%llu\n", + "graphs=%llu graph-hits=%llu graph-evictions=%llu launches=%llu " + "fused-decode-launches=%llu fused-decode-experts=%llu\n", runtime_->io->effective_backend_name(), (unsigned long long) stats.requests, (unsigned long long) stats.read_ops, @@ -996,7 +1239,9 @@ void MoeHybridStreamEngine::destroy() { (unsigned long long) runtime_->compute_stats.graph_builds, (unsigned long long) runtime_->compute_stats.graph_cache_hits, (unsigned long long) runtime_->compute_stats.graph_evictions, - (unsigned long long) runtime_->compute_stats.graph_launches); + (unsigned long long) runtime_->compute_stats.graph_launches, + (unsigned long long) runtime_->compute_stats.fused_decode_launches, + (unsigned long long) runtime_->compute_stats.fused_decode_experts); } runtime_->io->destroy(); } @@ -1908,6 +2153,160 @@ bool eval_moe_streamed_experts( engine.request_experts(batch.layer, unique_experts.data(), (int) unique_experts.size(), MoeNvmePriority::Demand); + + // Decode is latency-sensitive and normally selects several experts for a + // single token. If every selected expert is already device-resident, keep + // the complete fork/join on the GPU. A miss retains the pipelined path so + // expert N compute can overlap the upload of N+1. This removes the hot-path + // host boundary without changing routing, formats, or prefill behavior. + bool all_selected_resident = runtime.config.fused_decode && + batch.n_tokens == 1 && unique_experts.size() > 1; + if (all_selected_resident) { + for (const int32_t expert : unique_experts) { + const auto found = runtime.device_index.find( + device_key(batch.layer, expert)); + if (found == runtime.device_index.end() || found->second < 0 || + found->second >= (int) runtime.device_slots.size()) { + all_selected_resident = false; + break; + } + const auto & slot = + runtime.device_slots[(size_t) found->second]; + if (!slot.valid || !slot.cache_managed || + slot.compute_users != 0 || slot.key.layer != batch.layer || + slot.key.expert != expert) { + all_selected_resident = false; + break; + } + } + } + if (all_selected_resident) { + struct ActiveSlotSet { + MoeHybridStreamEngine & engine; + std::vector slots; + ~ActiveSlotSet() { + for (auto it = slots.rbegin(); it != slots.rend(); ++it) { + engine.release_device_slot(*it); + } + } + } active{engine, {}}; + active.slots.reserve(unique_experts.size()); + + std::vector bindings; + std::vector route_weights; + bindings.reserve(unique_experts.size()); + route_weights.reserve(unique_experts.size()); + for (const int32_t expert : unique_experts) { + float combined_weight = 0.0f; + for (int rank = 0; rank < batch.top_k; ++rank) { + if (batch.selected_ids[(size_t) rank] == expert) { + combined_weight += batch.selected_weights[(size_t) rank]; + } + } + if (!std::isfinite(combined_weight)) { + if (err) *err = "combined expert route weight overflowed"; + return false; + } + + int slot_index = -1; + if (!engine.stage_expert_cached_async( + batch.layer, expert, &slot_index, err) || + !engine.activate_device_slot(slot_index, err)) { + return false; + } + active.slots.push_back(slot_index); + const auto & slot = runtime.device_slots[(size_t) slot_index]; + if (!validate_moe_stream_expert_layout(spec, slot.layout, err)) { + return false; + } + const MoeExpertComponentKind gate_kind = spec.fused_gate_up + ? MoeExpertComponentKind::FusedGateUp + : MoeExpertComponentKind::Gate; + const auto * gate_component = + slot.device_layout.component(gate_kind); + const auto * up_component = spec.fused_gate_up + ? nullptr + : slot.device_layout.component(MoeExpertComponentKind::Up); + const auto * down_component = + slot.device_layout.component(MoeExpertComponentKind::Down); + if (!gate_component || !down_component || + (!spec.fused_gate_up && !up_component)) { + if (err) *err = + "streamed device layout is missing an expert component"; + return false; + } + const auto * base = static_cast(slot.data); + bindings.push_back({ + base + gate_component->offset, + up_component ? base + up_component->offset : nullptr, + base + down_component->offset, + gate_component->alloc_bytes, + up_component ? up_component->alloc_bytes : 0, + down_component->alloc_bytes, + }); + route_weights.push_back(combined_weight); + } + + std::unique_ptr ephemeral; + PersistentStreamMoEDecodeGraph * graph = nullptr; + const uint64_t touch = ++runtime.graph_clock; + if (runtime.config.graph_cache_entries > 0) { + for (auto & candidate : runtime.fused_decode_graph_cache) { + if (candidate && candidate->matches( + spec, (int) bindings.size())) { + candidate->last_touch = touch; + ++runtime.compute_stats.graph_cache_hits; + graph = candidate.get(); + break; + } + } + } + if (!graph) { + std::unique_ptr built( + new (std::nothrow) PersistentStreamMoEDecodeGraph); + if (!built) { + if (err) *err = + "failed to allocate fused streamed-MoE graph"; + return false; + } + if (!built->build(runtime.backend, runtime.device_pool_buffer, + spec, bindings, err)) { + return false; + } + built->last_touch = touch; + ++runtime.compute_stats.graph_builds; + if (runtime.config.graph_cache_entries <= 0) { + graph = built.get(); + ephemeral = std::move(built); + } else { + if ((int) runtime.fused_decode_graph_cache.size() >= + runtime.config.graph_cache_entries) { + auto victim = std::min_element( + runtime.fused_decode_graph_cache.begin(), + runtime.fused_decode_graph_cache.end(), + [](const auto & a, const auto & b) { + return a->last_touch < b->last_touch; + }); + if (victim != runtime.fused_decode_graph_cache.end()) { + runtime.fused_decode_graph_cache.erase(victim); + ++runtime.compute_stats.graph_evictions; + } + } + graph = built.get(); + runtime.fused_decode_graph_cache.push_back(std::move(built)); + } + } + if (!graph->launch( + bindings, batch.inputs, route_weights.data(), err) || + !graph->finish(out, err)) { + return false; + } + ++runtime.compute_stats.graph_launches; + ++runtime.compute_stats.fused_decode_launches; + runtime.compute_stats.fused_decode_experts += bindings.size(); + return true; + } + int staged_slot = -1; if (!engine.stage_expert_cached_async( batch.layer, unique_experts[0], &staged_slot, err)) { diff --git a/server/src/common/moe_hybrid_stream.h b/server/src/common/moe_hybrid_stream.h index 9eee5169c..f3e9b32c2 100644 --- a/server/src/common/moe_hybrid_stream.h +++ b/server/src/common/moe_hybrid_stream.h @@ -29,6 +29,11 @@ struct MoeStreamConfig { // activation, scales, and batch width. A small bounded cache removes graph // construction from decode without assuming every layer uses one format. int graph_cache_entries = 8; + // Decode normally routes one token to several experts. When that complete + // route set is device-resident, submit its independent branches as one + // backend graph and reduce on the GPU, avoiding one synchronization and + // D2H copy per expert. Misses and multi-token prefill retain the pipeline. + bool fused_decode = true; // Optional adaptive GPU expert-cache budget. Zero keeps only the pipeline // slots. The hardware planner can safely assign otherwise-unused Strix // memory here while retaining its KV/graph reserve. @@ -100,6 +105,8 @@ struct MoeStreamComputeStats { uint64_t graph_cache_hits = 0; uint64_t graph_evictions = 0; uint64_t graph_launches = 0; + uint64_t fused_decode_launches = 0; + uint64_t fused_decode_experts = 0; }; struct MoeStreamCacheWarmEntry { diff --git a/server/test/bench_kimi_k3_hetero.cpp b/server/test/bench_kimi_k3_hetero.cpp index 62dabb0d4..a1269cc65 100644 --- a/server/test/bench_kimi_k3_hetero.cpp +++ b/server/test/bench_kimi_k3_hetero.cpp @@ -478,12 +478,15 @@ int main(int argc, char ** argv) { "ssd_payload_gib=%.6f physical_gib=%.6f pipeline_gib_s=%.6f " "estimated_device_cache_hit=%.4f cache_gib=%.3f io_errors=%" PRIu64 " graph_builds=%" PRIu64 " graph_hits=%" PRIu64 - " graph_launches=%" PRIu64 "\n", + " graph_launches=%" PRIu64 " fused_decode_launches=%" PRIu64 + " fused_decode_experts=%" PRIu64 "\n", gib(stats.payload_bytes), gib(stats.physical_bytes), seconds > 0 ? gib(stats.payload_bytes) / seconds : 0.0, hit_rate, gib(engine.device_cache_bytes()), stats.errors, compute_stats.graph_builds, compute_stats.graph_cache_hits, - compute_stats.graph_launches); + compute_stats.graph_launches, + compute_stats.fused_decode_launches, + compute_stats.fused_decode_experts); engine.destroy(); ggml_backend_free(backend); diff --git a/server/test/test_moe_stream_compute.cpp b/server/test/test_moe_stream_compute.cpp index 7b3a2ad48..cb7185988 100644 --- a/server/test/test_moe_stream_compute.cpp +++ b/server/test/test_moe_stream_compute.cpp @@ -27,7 +27,7 @@ namespace { struct MoeStreamComputeFixture {}; -constexpr int kExperts = 3; +constexpr int kExperts = 4; // 256 deliberately does not satisfy CUDA/HIP's 512-element quantized matrix // row padding. The MXFP4 case below therefore exercises the padded GPU-slot // path that real Kimi-K3 exposed. @@ -239,17 +239,19 @@ std::vector cpu_reference( const std::vector & down, const std::vector & input, const int32_t * ids, - const float * weights) { + const float * weights, + int n_tokens = kTokens, + int top_k = kTopK) { constexpr float gate_scale = 0.8f; constexpr float up_scale = 1.1f; constexpr float down_scale = 0.9f; constexpr float beta = 4.0f; constexpr float linear_beta = 25.0f; - std::vector output((size_t) kTokens * kOutput, 0.0f); + std::vector output((size_t) n_tokens * kOutput, 0.0f); std::vector activated(kFf); - for (int token = 0; token < kTokens; ++token) { - for (int rank = 0; rank < kTopK; ++rank) { - const int expert = ids[token * kTopK + rank]; + for (int token = 0; token < n_tokens; ++token) { + for (int rank = 0; rank < top_k; ++rank) { + const int expert = ids[token * top_k + rank]; for (int row = 0; row < kFf; ++row) { float g = 0.0f; float u = 0.0f; @@ -275,13 +277,146 @@ std::vector cpu_reference( value += down[wi] * activated[(size_t) column]; } output[(size_t) token * kOutput + row] += - weights[token * kTopK + rank] * down_scale * value; + weights[token * top_k + rank] * down_scale * value; } } } return output; } +void run_fused_decode_case(ggml_backend_t backend, bool mxfp4) { + std::vector gate; + std::vector up; + std::vector down; + fill_weights(gate, up, down); + std::vector gate_reference = gate; + std::vector up_reference = up; + std::vector down_reference = down; + ModelBytes model = mxfp4 + ? make_mxfp4_model_bytes( + gate, up, down, gate_reference, up_reference, down_reference) + : make_model_bytes(true, gate, up, down); + TempFile file(model.file); + + MoeHybridStorage storage; + storage.mmap_size = model.file.size(); + storage.mmap_fd = ::dup(file.fd); + STREAM_REQUIRE(storage.mmap_fd >= 0); + storage.layer_regions.push_back(model.regions); + + MoeStreamConfig config; + config.device_slots = kExperts; + config.device_cache_bytes = 0; + config.graph_cache_entries = 4; + config.fused_decode = true; + config.nvme.backend = MoeNvmeBackend::ThreadPool; + config.nvme.direct_io = MoeNvmeDirectMode::Disabled; + config.nvme.host_slots = 6; + config.nvme.io_threads = 2; + + MoeHybridStreamEngine engine; + std::string error; + STREAM_REQUIRE(engine.init( + backend, model.slot_bytes, storage, config, &error)); + + MoeStreamExpertSpec spec; + spec.input_dim = kInput; + spec.intermediate_dim = kFf; + spec.output_dim = kOutput; + spec.gate_type = mxfp4 ? GGML_TYPE_MXFP4 : GGML_TYPE_F32; + spec.up_type = mxfp4 ? GGML_TYPE_MXFP4 : GGML_TYPE_F32; + spec.down_type = mxfp4 ? GGML_TYPE_MXFP4 : GGML_TYPE_F32; + spec.gated_activation = MoeGatedActivation::Situ; + spec.gate_scale = 0.8f; + spec.up_scale = 1.1f; + spec.down_scale = 0.9f; + + std::vector input((size_t) kInput); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = 0.12f * std::sin(0.07f * (float) (i + 1)); + } + constexpr int kDecodeTopK = 3; + int32_t ids[kDecodeTopK] = {2, 0, 1}; + float weights[kDecodeTopK] = {0.50f, 0.30f, 0.20f}; + MoeStreamRouteBatch batch; + batch.layer = 0; + batch.n_expert = kExperts; + batch.top_k = kDecodeTopK; + batch.n_tokens = 1; + batch.inputs = input.data(); + batch.selected_ids = ids; + batch.selected_weights = weights; + + // Prepare the padded numerical layout without admitting a pinned entry. + MoeStreamCacheWarmStats prepare_stats; + STREAM_REQUIRE(engine.warm_and_pin_device_cache( + {spec}, {{0, 0, 1, model.slot_bytes}}, kExperts, + &prepare_stats, &error)); + STREAM_REQUIRE(prepare_stats.capacity_drops == 1); + + const std::vector expected = cpu_reference( + gate_reference, up_reference, down_reference, + input, ids, weights, 1, kDecodeTopK); + std::vector actual; + auto require_close = [&](const std::vector & reference) { + STREAM_REQUIRE(actual.size() == reference.size()); + for (size_t i = 0; i < actual.size(); ++i) { + const float tolerance = mxfp4 + ? 2.0e-4f + 2.0e-3f * std::fabs(reference[i]) + : 2.0e-5f + 2.0e-4f * std::fabs(reference[i]); + STREAM_REQUIRE(std::fabs(actual[i] - reference[i]) <= tolerance); + } + }; + + // A cold route must preserve the transfer/compute overlap pipeline. + STREAM_REQUIRE(eval_moe_streamed_experts( + engine, spec, batch, actual, &error)); + require_close(expected); + const MoeStreamComputeStats cold = engine.compute_stats(); + STREAM_REQUIRE(cold.graph_launches == kDecodeTopK); + STREAM_REQUIRE(cold.fused_decode_launches == 0); + + // Populate every slot, then verify the all-resident fused path. + for (int expert = 0; expert < kExperts; ++expert) { + int slot = -1; + STREAM_REQUIRE(engine.stage_expert_cached_async( + 0, expert, &slot, &error)); + STREAM_REQUIRE(engine.activate_device_slot(slot, &error)); + engine.release_device_slot(slot); + } + STREAM_REQUIRE(eval_moe_streamed_experts( + engine, spec, batch, actual, &error)); + require_close(expected); + + const MoeStreamComputeStats first = engine.compute_stats(); + STREAM_REQUIRE(first.graph_builds == cold.graph_builds + 1); + STREAM_REQUIRE(first.graph_launches == cold.graph_launches + 1); + STREAM_REQUIRE(first.fused_decode_launches == 1); + STREAM_REQUIRE(first.fused_decode_experts == kDecodeTopK); + + // Reuse the same graph shape with a different expert set and ordering. + // This catches stale captured device pointers in CUDA/HIP graph mode. + ids[0] = 3; + ids[1] = 2; + ids[2] = 0; + weights[0] = 0.25f; + weights[1] = 0.60f; + weights[2] = 0.15f; + const std::vector expected_rebound = cpu_reference( + gate_reference, up_reference, down_reference, + input, ids, weights, 1, kDecodeTopK); + STREAM_REQUIRE(eval_moe_streamed_experts( + engine, spec, batch, actual, &error)); + require_close(expected_rebound); + const MoeStreamComputeStats second = engine.compute_stats(); + STREAM_REQUIRE(second.graph_builds == first.graph_builds); + STREAM_REQUIRE(second.graph_cache_hits > first.graph_cache_hits); + STREAM_REQUIRE(second.graph_launches == cold.graph_launches + 2); + STREAM_REQUIRE(second.fused_decode_launches == 2); + STREAM_REQUIRE(second.fused_decode_experts == 2 * kDecodeTopK); + engine.destroy(); +} + void run_layout_case(ggml_backend_t backend, bool expert_major) { std::vector gate; std::vector up; @@ -524,6 +659,8 @@ TEST_CASE(MoeStreamComputeFixture, persistent_graph_matches_cpu_and_padded_mxfp4 } run_layout_case(backend, false); run_layout_case(backend, true); + run_fused_decode_case(backend, false); + run_fused_decode_case(backend, true); run_mxfp4_padding_case(backend); run_pinned_cache_case(backend); ggml_backend_free(backend); From 90b126c1bda792a04c5110b2d676e4eed6ebd719 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:02:27 +0200 Subject: [PATCH 16/20] test(kimi): record full-model Strix qualification --- server/docs/KIMI_K3_HETERO.md | 51 ++++++++++++++----- .../scripts/benchmark_kimi_k3_deployments.py | 50 +++++++++++------- .../test_benchmark_kimi_k3_deployments.py | 17 +++++++ 3 files changed, 88 insertions(+), 30 deletions(-) diff --git a/server/docs/KIMI_K3_HETERO.md b/server/docs/KIMI_K3_HETERO.md index d0f196a5c..ee7912ba2 100644 --- a/server/docs/KIMI_K3_HETERO.md +++ b/server/docs/KIMI_K3_HETERO.md @@ -154,6 +154,34 @@ oracle on both backends and also change expert addresses and route weights between cached graph launches. These numbers isolate the routed core and do not claim complete-model token speed. +## Full 2.8T model qualification, 2026-08-03 + +The complete 14-shard Unsloth IQ1_S checkpoint was run through the native text +backend on Lucebox4's Strix Halo. This was an end-to-end greedy generation, +not synthetic byte ranges or a routed-core microbenchmark. The run used commit +`3450da3f`, `--moe-storage ssd`, automatic cache sizing, direct `io_uring`, a +512-token context allocation, and disabled HTTP/prefix caches. + +| Result | Cold request | Warm request | +|---|---:|---:| +| Startup | 46.0 s | same server | +| 34-token prefill | 129.2 s | 103.6 s | +| Four-token decode | 16.0 s | 16.3 s | +| Decode rate | **0.200 token/s** | **0.200 token/s** | + +Both requests returned the identical greedy prefix `Hello! Nice to`. The +loader placed 57.93 GiB of non-routed weights on Strix and assigned 32.66 GiB +of the 495.26 GiB routed pool to its adaptive cache. Across both requests the +engine read 262.31 GiB at 1.158 GiB/s, obtained a 7.5% aggregate cache-hit +rate, issued 86 fused resident-decode launches, and reported zero I/O errors +or timeouts. + +This establishes that the released 2.8T checkpoint fits and generates on a +128 GiB Strix-only system. It also establishes the current bottleneck: real +token-sequential prefill plus the small cache yields about one decoded token +every five seconds. Fused warm graphs cannot materially change that result +while more than 92% of routed expert accesses miss the device cache. + ## Practical Lucebox placement The machine has about 125.08 GiB of system/UMA memory plus 31.86 GiB on the @@ -176,23 +204,20 @@ upload path measured below Strix. After approximately 57.94 GiB of non-routed weights plus OS, workspace, and a moderate context reserve, roughly 70-85 GiB may remain for routed experts. Under a uniform balanced-routing assumption this covers about 14-17% of the -routed pool. At the measured 3.804 GiB/s, the storage-only ceiling is then -approximately 0.50-0.51 token/s. Full Strix-only inference will be lower -because dense/recurrent work shares the same device; a later R9700 split may -recover part of that gap through overlap. - -So the honest expectation for this quant is **roughly one token every two to -three seconds**, not interactive multi-token-per-second generation. Real -router locality can move that estimate; only a route trace from the real model -can establish it. +routed pool. The earlier routed-core storage ceiling was 0.50-0.51 token/s, +but full Strix-only inference shares the device with dense and recurrent work. +The full checkpoint measured **0.200 token/s**, or roughly one decoded token +every five seconds. Real route traces and heterogeneous overlap are required +to determine how much of that gap is recoverable. -## Next full-model milestone +## Next optimization milestone The implementation no longer needs another generic cache, a second Kimi -router, or per-layer worker creation. Full-scale qualification requires: +router, per-layer worker creation, or a first full-scale smoke test. The next +work is: -1. Stage all 14 IQ1_S shards and run a short token-for-token comparison against - the upstream Kimi implementation. +1. Compare a longer deterministic sample token-for-token against the upstream + Kimi implementation. 2. Record real `(layer, expert)` routes on a calibration prompt suite and let the existing placement planner allocate the measured best cache under the chosen context budget. diff --git a/server/scripts/benchmark_kimi_k3_deployments.py b/server/scripts/benchmark_kimi_k3_deployments.py index 1e5dbd757..974124187 100644 --- a/server/scripts/benchmark_kimi_k3_deployments.py +++ b/server/scripts/benchmark_kimi_k3_deployments.py @@ -28,17 +28,20 @@ _SPLIT_GGUF = re.compile(r"^(?P.+)-(?P\d+)-of-(?P\d+)(?P\.gguf)$") -_NVME_TELEMETRY = re.compile( - r"\[moe-nvme\] io=(?P\S+) requests=(?P\d+) reads=(?P\d+) " - r"payload=(?P[\d.]+) GiB physical=(?P[\d.]+) GiB " - r"active-io-rate=(?P[\d.]+) GiB/s cache-hit=(?P[\d.]+)% " - r"mean-demand-wait=(?P[\d.]+) ms .*?timeouts=(?P\d+) " - r"errors=(?P\d+) device-cache=(?P[\d.]+) MiB " - r"slots=(?P\d+) hits=(?P\d+) misses=(?P\d+) " - r"evictions=(?P\d+) graphs=(?P\d+) " - r"graph-hits=(?P\d+) graph-evictions=(?P\d+) " - r"launches=(?P\d+)" +_NVME_KEY_VALUE = re.compile( + r"(?[a-z][a-z0-9-]*)=(?P\S+)" ) +_NVME_RESULT_NAMES = { + "payload": "payload_gib", + "physical": "physical_gib", + "active-io-rate": "active_io_gib_s", + "cache-hit": "cache_hit_pct", + "mean-demand-wait": "mean_demand_wait_ms", + "device-cache": "device_cache_mib", + "hits": "device_hits", + "misses": "device_misses", + "evictions": "device_evictions", +} @dataclass(frozen=True) @@ -175,13 +178,26 @@ def extract_nvme_telemetry(log_path: Path) -> list[dict[str, Any]]: telemetry: list[dict[str, Any]] = [] if not log_path.exists(): return telemetry - for match in _NVME_TELEMETRY.finditer(log_path.read_text(errors="replace")): - row: dict[str, Any] = {"io": match.group("io")} - for key, value in match.groupdict().items(): - if key == "io": - continue - row[key] = float(value) if "." in value else int(value) - telemetry.append(row) + for line in log_path.read_text(errors="replace").splitlines(): + if not line.startswith("[moe-nvme] "): + continue + row: dict[str, Any] = {} + for match in _NVME_KEY_VALUE.finditer(line): + raw_key = match.group("key") + key = _NVME_RESULT_NAMES.get(raw_key, raw_key.replace("-", "_")) + raw_value = match.group("value") + numeric_value = raw_value.removesuffix("%") + try: + value: Any = ( + float(numeric_value) + if "." in numeric_value + else int(numeric_value) + ) + except ValueError: + value = raw_value + row[key] = value + if "io" in row: + telemetry.append(row) return telemetry diff --git a/server/scripts/test_benchmark_kimi_k3_deployments.py b/server/scripts/test_benchmark_kimi_k3_deployments.py index 6162115be..2cf278d73 100644 --- a/server/scripts/test_benchmark_kimi_k3_deployments.py +++ b/server/scripts/test_benchmark_kimi_k3_deployments.py @@ -104,6 +104,23 @@ def test_extracts_each_owner_telemetry_line(self): self.assertEqual(rows[0]["active_io_gib_s"], 3.75) self.assertEqual(rows[0]["errors"], 0) + def test_telemetry_parser_tolerates_new_and_reordered_fields(self): + line = ( + "[moe-nvme] io=io_uring+direct launches=107638 requests=87143 " + "payload=261.811 GiB cache-hit=7.5% device-cache=33444.1 MiB " + "slots=5436 pinned=0 hits=65354 misses=43574 evictions=38138 " + "fused-decode-launches=86 fused-decode-experts=1376 future-counter=7\n" + ) + with tempfile.TemporaryDirectory() as directory: + log = Path(directory) / "server.log" + log.write_text(line) + rows = extract_nvme_telemetry(log) + self.assertEqual(rows[0]["io"], "io_uring+direct") + self.assertEqual(rows[0]["device_hits"], 65354) + self.assertEqual(rows[0]["pinned"], 0) + self.assertEqual(rows[0]["fused_decode_launches"], 86) + self.assertEqual(rows[0]["future_counter"], 7) + if __name__ == "__main__": unittest.main() From 48329260f3b51ee563385b4d00c4743cfc2f74a3 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 3 Aug 2026 23:31:49 +0200 Subject: [PATCH 17/20] feat(kimi): stream experts across mixed HIP CUDA owners --- server/docs/MOE_NVME_STREAMING.md | 37 ++- server/src/common/moe_hybrid_stream.cpp | 332 ++++++++++++--------- server/src/deepseek4/deepseek4_backend.cpp | 2 +- server/src/kimi_k3/kimi_k3_backend.cpp | 59 +++- server/src/kimi_k3/kimi_k3_backend.h | 10 +- server/test/test_moe_stream_compute.cpp | 21 +- 6 files changed, 284 insertions(+), 177 deletions(-) diff --git a/server/docs/MOE_NVME_STREAMING.md b/server/docs/MOE_NVME_STREAMING.md index c2dfebb44..62320a784 100644 --- a/server/docs/MOE_NVME_STREAMING.md +++ b/server/docs/MOE_NVME_STREAMING.md @@ -218,10 +218,10 @@ export DFLASH_MOE_HOTNESS_CSV=/models/kimi-routes.csv ``` At least one quarter of each device cache (and never fewer than two slots) -remains adaptive for profile drift and misses. With two GPUs, the same runtime -ownership rule filters both warm plans, so an expert is never pinned on both -R9700 and Strix. A missing profile simply retains the adaptive LFRU cache; -prediction is not required for correctness. +remains adaptive for profile drift and misses. With two GPU owners, the same +deterministic ownership rule filters both warm plans, so an expert is never +pinned on both devices. A missing profile simply retains the adaptive LFRU +cache; prediction is not required for correctness. On a two-GPU Lucebox, put the compute-intensive primary path on the R9700 and use Strix as the secondary capacity owner. Both devices receive independent @@ -237,13 +237,26 @@ export DFLASH_MOE_TP_GPU= This is functional expert ownership, not a contiguous layer split, so do not use `--target-devices`. `DFLASH_MOE_PLACEMENT` may point at an offline -`MoeHybridPlacement` JSON; its hot expert IDs become R9700-owned and -all other selected routes become Strix-owned. Without a plan, +`MoeHybridPlacement` JSON; its hot expert IDs become R9700-owned and all other +selected routes become Strix-owned. Without a plan, `DFLASH_MOE_PRIMARY_SHARE_PER_MILLE` controls a deterministic bring-up split -(default 500). Leaving `DFLASH_MOE_TP_GPU` unset preserves the single-device -path, including Strix-only systems. The current Kimi join still uses -activation-sized host staging; a device-resident peer join remains the next -throughput optimization. +(default 500). Leaving both secondary-owner variables unset preserves the +single-device path, including Strix-only systems. + +Mixed-vendor builds use the same path. For example, a HIP primary with an RTX +secondary selects the isolated CUDA module in the existing process: + +```bash +export DFLASH_MOE_TP_BACKEND=cuda +export DFLASH_MOE_TP_GPU=0 + +./build-hip-mixed/dflash_server \ + /path/to/Kimi-K3-UD-IQ1_S-00001-of-00014.gguf \ + --target-device hip:0 --max-ctx 8192 +``` + +The current Kimi join still uses activation-sized host staging; a +device-resident peer join remains the next throughput optimization. ## Tuning and diagnostics @@ -272,6 +285,7 @@ not improve the qualified P310 drive and consumes extra pinned/system memory. | `DFLASH_MOE_ROUTE_STATS_OUT` | unset | Record Kimi native-router counts as a reusable model-neutral CSV profile | | `DFLASH_MOE_HOTNESS_CSV` | unset | Warm and pin the highest-value experts from a compatible routing profile | | `DFLASH_MOE_TP_GPU` | primary GPU | Optional secondary GPU; enables concurrent route ownership when different from the primary | +| `DFLASH_MOE_TP_BACKEND` | primary backend | Optional secondary runtime (`cuda` or `hip`); a different runtime enables mixed-vendor route ownership even when both use device index `0` | | `DFLASH_MOE_PLACEMENT` | unset | Offline placement JSON; listed experts belong to the primary GPU | | `DFLASH_MOE_PRIMARY_SHARE_PER_MILLE` | `500` | Bring-up hash split used only when no placement is supplied | | `DFLASH_MOE_DUAL_STREAM_TRACE` | unset | Debug per-layer owner counts and branch/wall timing | @@ -285,7 +299,8 @@ only when it covers the complete logical payload at an unaligned file tail. `test_moe_stream_compute` generates tiny experts and checks both tensor-major and expert-major GPU results against a CPU oracle. It defaults to GPU 0, so it runs directly on Strix-only systems; -`DFLASH_TEST_GPU` selects another device on multi-GPU hosts. The standalone +`DFLASH_TEST_GPU` selects another device and `DFLASH_TEST_BACKEND` selects the +linked or dynamically loaded runtime on mixed builds. The standalone targets `test_moe_nvme_scheduler`, `test_moe_expert_package`, `bench_moe_nvme_io`, and `bench_moe_nvme_pipeline` test scheduling, raw storage, and the complete diff --git a/server/src/common/moe_hybrid_stream.cpp b/server/src/common/moe_hybrid_stream.cpp index 3d3af5bd3..08e456e2a 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -1,9 +1,7 @@ #include "moe_hybrid_stream.h" -#include "gpu_runtime_compat.h" #include "ggml-alloc.h" #include "ggml-backend.h" -#include "ggml-cuda.h" #include #include @@ -29,13 +27,93 @@ namespace dflash::common { namespace { -bool pinned_allocate(void ** ptr, size_t bytes, void *) { - return cudaMallocHost(ptr, bytes) == cudaSuccess; -} +// Allocate host staging through the selected ggml backend module. This is +// essential in a mixed HIP+CUDA process: memory pinned by the linked HIP +// runtime is not CUDA-pinned memory for a dynamically loaded CUDA backend. +// Keeping allocation behind the backend device also makes the SSD scheduler +// independent of which GPU vendor owns a route partition. +class BackendHostAllocator { +public: + bool init(ggml_backend_t backend, std::string * err) { + const ggml_backend_dev_t device = backend + ? ggml_backend_get_device(backend) : nullptr; + buffer_type_ = device + ? ggml_backend_dev_host_buffer_type(device) : nullptr; + if (!buffer_type_) buffer_type_ = ggml_backend_cpu_buffer_type(); + if (!buffer_type_) { + if (err) *err = "stream backend has no host staging buffer type"; + return false; + } + return true; + } -void pinned_free(void * ptr, void *) { - if (ptr) (void) cudaFreeHost(ptr); -} + bool allocate(void ** ptr, size_t bytes) { + if (!ptr || bytes == 0 || !buffer_type_) return false; + *ptr = nullptr; + ggml_backend_buffer_t buffer = + ggml_backend_buft_alloc_buffer(buffer_type_, bytes); + if (!buffer) return false; + void * base = ggml_backend_buffer_get_base(buffer); + if (!base) { + ggml_backend_buffer_free(buffer); + return false; + } + try { + std::lock_guard lock(mutex_); + const auto [_, inserted] = buffers_.emplace(base, buffer); + if (!inserted) { + ggml_backend_buffer_free(buffer); + return false; + } + } catch (const std::bad_alloc &) { + ggml_backend_buffer_free(buffer); + return false; + } + *ptr = base; + return true; + } + + void release(void * ptr) { + if (!ptr) return; + ggml_backend_buffer_t buffer = nullptr; + { + std::lock_guard lock(mutex_); + const auto found = buffers_.find(ptr); + if (found == buffers_.end()) return; + buffer = found->second; + buffers_.erase(found); + } + ggml_backend_buffer_free(buffer); + } + + static bool allocate_callback(void ** ptr, size_t bytes, void * opaque) { + return opaque && static_cast(opaque)->allocate( + ptr, bytes); + } + + static void free_callback(void * ptr, void * opaque) { + if (opaque) static_cast(opaque)->release(ptr); + } + + ~BackendHostAllocator() { + for (;;) { + ggml_backend_buffer_t buffer = nullptr; + { + std::lock_guard lock(mutex_); + if (buffers_.empty()) break; + const auto found = buffers_.begin(); + buffer = found->second; + buffers_.erase(found); + } + ggml_backend_buffer_free(buffer); + } + } + +private: + ggml_backend_buffer_type_t buffer_type_ = nullptr; + std::mutex mutex_; + std::unordered_map buffers_; +}; int env_bounded_int(const char * name, int fallback, int lo, int hi) { const char * value = std::getenv(name); @@ -100,41 +178,6 @@ uint64_t device_key(int layer, int expert) { return ((uint64_t) (uint32_t) layer << 32) | (uint32_t) expert; } -int backend_device_index(ggml_backend_t backend) { - if (!backend || !ggml_backend_is_cuda(backend)) return -1; - ggml_backend_dev_t wanted = ggml_backend_get_device(backend); - ggml_backend_reg_t reg = ggml_backend_cuda_reg(); - const int count = ggml_backend_cuda_get_device_count(); - for (int device = 0; device < count; ++device) { - if (ggml_backend_reg_dev_get(reg, (size_t) device) == wanted) return device; - } - return -1; -} - -// HIP/CUDA streams, events, and allocations belong to the current device. -// The heterogeneous engine alternates R9700 and Strix backends on one host -// thread, so relying on whichever backend ran last is a cross-device bug. -class ScopedGpuDevice { -public: - explicit ScopedGpuDevice(int target) : target_(target) { - if (target_ < 0 || cudaGetDevice(&previous_) != cudaSuccess) return; - valid_ = true; - if (previous_ != target_) switched_ = cudaSetDevice(target_) == cudaSuccess; - } - - ~ScopedGpuDevice() { - if (valid_ && switched_) (void) cudaSetDevice(previous_); - } - - bool ready() const { return valid_ && (previous_ == target_ || switched_); } - -private: - int target_ = -1; - int previous_ = -1; - bool valid_ = false; - bool switched_ = false; -}; - } // namespace MoeStreamConfig MoeStreamConfig::from_env() { @@ -801,7 +844,8 @@ struct MoeHybridStreamEngine::Runtime { struct DeviceSlot { void * data = nullptr; - cudaEvent_t ready = nullptr; + ggml_tensor * transfer_tensor = nullptr; + ggml_backend_event_t ready = nullptr; bool pending = false; bool valid = false; bool cache_managed = false; @@ -816,13 +860,17 @@ struct MoeHybridStreamEngine::Runtime { }; ggml_backend_t backend = nullptr; - int device = -1; + // A second backend instance on the same device owns the upload stream. + // Its interface comes from the same module as `backend`, so this works for + // both the linked runtime and an isolated CUDA/HIP peer module. + ggml_backend_t transfer_backend = nullptr; size_t max_expert_bytes = 0; MoeStreamConfig config{}; + BackendHostAllocator host_allocator; std::unique_ptr io; - cudaStream_t transfer_stream = nullptr; ggml_backend_buffer_t device_pool_buffer = nullptr; void * device_pool = nullptr; + ggml_context * device_slot_ctx = nullptr; size_t device_stride = 0; size_t device_pool_bytes = 0; std::vector device_slots; @@ -845,16 +893,17 @@ struct MoeHybridStreamEngine::Runtime { template void release_device_cache(RuntimeT & runtime) { if (runtime.backend) ggml_backend_synchronize(runtime.backend); - if (runtime.transfer_stream) { - (void) cudaStreamSynchronize(runtime.transfer_stream); + if (runtime.transfer_backend) { + ggml_backend_synchronize(runtime.transfer_backend); } runtime.graph_cache.clear(); runtime.fused_decode_graph_cache.clear(); for (auto & slot : runtime.device_slots) { slot.host_lease.reset(); - if (slot.ready) (void) cudaEventDestroy(slot.ready); + if (slot.ready) ggml_backend_event_free(slot.ready); slot.ready = nullptr; slot.data = nullptr; + slot.transfer_tensor = nullptr; slot.pending = false; } runtime.device_slots.clear(); @@ -866,6 +915,8 @@ void release_device_cache(RuntimeT & runtime) { } runtime.device_pool_buffer = nullptr; runtime.device_pool = nullptr; + if (runtime.device_slot_ctx) ggml_free(runtime.device_slot_ctx); + runtime.device_slot_ctx = nullptr; runtime.device_stride = 0; runtime.device_pool_bytes = 0; } @@ -928,8 +979,46 @@ bool allocate_device_cache(RuntimeT & runtime, std::string * err, return false; } auto * base = static_cast(runtime.device_pool); + if (runtime.device_stride > + static_cast(std::numeric_limits::max()) || + attempt_slots > + (std::numeric_limits::max() - 1024) / + ggml_tensor_overhead()) { + if (err) *err = "SSD GPU cache tensor size overflow"; + release_device_cache(runtime); + return false; + } + ggml_init_params params{}; + params.mem_size = attempt_slots * ggml_tensor_overhead() + 1024; + params.no_alloc = true; + runtime.device_slot_ctx = ggml_init(params); + if (!runtime.device_slot_ctx) { + if (err) *err = "failed to allocate SSD GPU cache tensor metadata"; + release_device_cache(runtime); + return false; + } + ggml_backend_buffer_clear(runtime.device_pool_buffer, 0); + // The pool is shared by the compute and upload backend instances. Finish + // initialization on the compute stream before the upload stream can reuse + // any slot, otherwise zero-filled quantization padding could race H2D. + ggml_backend_synchronize(runtime.backend); for (size_t i = 0; i < runtime.device_slots.size(); ++i) { - runtime.device_slots[i].data = base + i * runtime.device_stride; + auto & slot = runtime.device_slots[i]; + slot.data = base + i * runtime.device_stride; + slot.transfer_tensor = ggml_new_tensor_1d( + runtime.device_slot_ctx, GGML_TYPE_I8, + (int64_t) runtime.device_stride); + if (!slot.transfer_tensor || + ggml_backend_buffer_get_alloc_size( + runtime.device_pool_buffer, slot.transfer_tensor) > + runtime.device_stride || + ggml_backend_tensor_alloc( + runtime.device_pool_buffer, slot.transfer_tensor, + slot.data) != GGML_STATUS_SUCCESS) { + if (err) *err = "failed to bind SSD GPU cache transfer tensor"; + release_device_cache(runtime); + return false; + } } return true; } @@ -1084,10 +1173,19 @@ bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, return false; } runtime->backend = gpu_backend; - runtime->device = backend_device_index(gpu_backend); - ScopedGpuDevice device_scope(runtime->device); - if (!device_scope.ready()) { - if (err) *err = "failed to resolve/select SSD stream GPU"; + ggml_backend_dev_t device = ggml_backend_get_device(gpu_backend); + if (!device) { + if (err) *err = "failed to resolve SSD stream backend device"; + return false; + } + runtime->transfer_backend = ggml_backend_dev_init(device, nullptr); + if (!runtime->transfer_backend) { + if (err) *err = "failed to create SSD upload backend stream"; + return false; + } + if (!runtime->host_allocator.init(runtime->transfer_backend, err)) { + ggml_backend_free(runtime->transfer_backend); + runtime->transfer_backend = nullptr; return false; } runtime->max_expert_bytes = max_expert_bytes; @@ -1096,17 +1194,17 @@ bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, runtime->io.reset(new (std::nothrow) MoeNvmeScheduler); if (!runtime->io) { if (err) *err = "failed to allocate SSD scheduler"; + ggml_backend_free(runtime->transfer_backend); + runtime->transfer_backend = nullptr; return false; } if (!runtime->io->init(runtime->config.nvme, max_expert_bytes, - pinned_allocate, pinned_free, nullptr, err)) { - return false; - } - - cudaError_t gpu_err = cudaStreamCreate(&runtime->transfer_stream); - if (gpu_err != cudaSuccess) { - if (err) *err = std::string("failed to create SSD transfer stream: ") + - cudaGetErrorString(gpu_err); + BackendHostAllocator::allocate_callback, + BackendHostAllocator::free_callback, + &runtime->host_allocator, err)) { + runtime->io.reset(); + ggml_backend_free(runtime->transfer_backend); + runtime->transfer_backend = nullptr; return false; } if (!allocate_device_cache(*runtime, err)) { @@ -1161,7 +1259,7 @@ bool MoeHybridStreamEngine::bind_sources( bool MoeHybridStreamEngine::is_ready() const { return runtime_ && runtime_->backend && runtime_->io && - runtime_->io->is_initialized() && runtime_->transfer_stream && + runtime_->io->is_initialized() && runtime_->transfer_backend && !runtime_->device_slots.empty(); } @@ -1173,18 +1271,18 @@ void MoeHybridStreamEngine::destroy() { if (!runtime_) return; const size_t device_cache_slot_count = runtime_->device_slots.size(); const size_t device_cache_byte_count = runtime_->device_pool_bytes; - ScopedGpuDevice device_scope(runtime_->device); if (runtime_->backend) ggml_backend_synchronize(runtime_->backend); - if (runtime_->transfer_stream) { - (void) cudaStreamSynchronize(runtime_->transfer_stream); + if (runtime_->transfer_backend) { + ggml_backend_synchronize(runtime_->transfer_backend); } runtime_->graph_cache.clear(); runtime_->fused_decode_graph_cache.clear(); for (Runtime::DeviceSlot & slot : runtime_->device_slots) { slot.host_lease.reset(); - if (slot.ready) (void) cudaEventDestroy(slot.ready); + if (slot.ready) ggml_backend_event_free(slot.ready); slot.ready = nullptr; slot.data = nullptr; + slot.transfer_tensor = nullptr; slot.pending = false; } runtime_->device_slots.clear(); @@ -1194,8 +1292,8 @@ void MoeHybridStreamEngine::destroy() { } runtime_->device_pool_buffer = nullptr; runtime_->device_pool = nullptr; - if (runtime_->transfer_stream) (void) cudaStreamDestroy(runtime_->transfer_stream); - runtime_->transfer_stream = nullptr; + if (runtime_->device_slot_ctx) ggml_free(runtime_->device_slot_ctx); + runtime_->device_slot_ctx = nullptr; if (runtime_->io) { const MoeNvmeStats stats = runtime_->io->stats(); if (stats.requests != 0 || stats.read_ops != 0 || stats.errors != 0) { @@ -1245,6 +1343,11 @@ void MoeHybridStreamEngine::destroy() { } runtime_->io->destroy(); } + runtime_->io.reset(); + if (runtime_->transfer_backend) { + ggml_backend_free(runtime_->transfer_backend); + runtime_->transfer_backend = nullptr; + } runtime_.reset(); } @@ -1314,11 +1417,6 @@ bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, if (err) *err = "stream engine has no bound SSD model source"; return false; } - ScopedGpuDevice device_scope(runtime_->device); - if (!device_scope.ready()) { - if (err) *err = "failed to select SSD stream GPU"; - return false; - } if (device_slot < 0 || device_slot >= (int) runtime_->device_slots.size()) { if (err) *err = "SSD device slot is out of range"; return false; @@ -1333,12 +1431,7 @@ bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, return false; } if (dst.pending) { - const cudaError_t wait_err = cudaEventSynchronize(dst.ready); - if (wait_err != cudaSuccess) { - if (err) *err = std::string("failed waiting for prior SSD upload: ") + - cudaGetErrorString(wait_err); - return false; - } + ggml_backend_event_synchronize(dst.ready); dst.pending = false; dst.host_lease.reset(); } @@ -1353,14 +1446,17 @@ bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, dst.device_layout = Runtime::DeviceExpertLayout{}; if (!dst.ready) { - const cudaError_t event_create_err = - cudaEventCreateWithFlags(&dst.ready, cudaEventDisableTiming); - if (event_create_err != cudaSuccess) { - if (err) *err = std::string("failed to create expert upload event: ") + - cudaGetErrorString(event_create_err); + dst.ready = ggml_backend_event_new( + ggml_backend_get_device(runtime_->transfer_backend)); + if (!dst.ready) { + if (err) *err = "stream backend does not support upload events"; return false; } } + if (!dst.transfer_tensor) { + if (err) *err = "SSD device slot has no transfer tensor"; + return false; + } MoeNvmeLease lease; if (!runtime_->io->acquire(layer, expert_id, lease, err)) return false; @@ -1406,26 +1502,10 @@ bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, return false; } - cudaError_t gpu_err = cudaMemcpyAsync( - static_cast(dst.data) + device_component.offset, - source, device_component.logical_bytes, - cudaMemcpyHostToDevice, runtime_->transfer_stream); - if (gpu_err == cudaSuccess && - device_component.alloc_bytes > device_component.logical_bytes) { - gpu_err = cudaMemsetAsync( - static_cast(dst.data) + device_component.offset + - device_component.logical_bytes, - 0, - device_component.alloc_bytes - - device_component.logical_bytes, - runtime_->transfer_stream); - } - if (gpu_err != cudaSuccess) { - (void) cudaStreamSynchronize(runtime_->transfer_stream); - if (err) *err = std::string("asynchronous expert H2D failed: ") + - cudaGetErrorString(gpu_err); - return false; - } + ggml_backend_tensor_set_async( + runtime_->transfer_backend, dst.transfer_tensor, + source, device_component.offset, + device_component.logical_bytes); } dst.device_layout = device_layout; } else { @@ -1434,16 +1514,10 @@ bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, // always registers an exact backend-padded layout before reaching here. for (int i = 0; i < lease.layout().span_count; ++i) { const MoeExpertIoSpan & span = lease.layout().spans[i]; - cudaError_t gpu_err = cudaMemcpyAsync( - static_cast(dst.data) + span.device_offset, + ggml_backend_tensor_set_async( + runtime_->transfer_backend, dst.transfer_tensor, lease.data() + span.buffer_offset, - span.bytes, cudaMemcpyHostToDevice, runtime_->transfer_stream); - if (gpu_err != cudaSuccess) { - (void) cudaStreamSynchronize(runtime_->transfer_stream); - if (err) *err = std::string("asynchronous expert H2D failed: ") + - cudaGetErrorString(gpu_err); - return false; - } + span.device_offset, span.bytes); } dst.device_layout.component_count = lease.layout().component_count; dst.device_layout.bytes = lease.layout().payload_bytes; @@ -1455,13 +1529,7 @@ bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, component.bytes, component.bytes}; } } - const cudaError_t event_err = cudaEventRecord(dst.ready, runtime_->transfer_stream); - if (event_err != cudaSuccess) { - (void) cudaStreamSynchronize(runtime_->transfer_stream); - if (err) *err = std::string("failed to record expert upload event: ") + - cudaGetErrorString(event_err); - return false; - } + ggml_backend_event_record(dst.ready, runtime_->transfer_backend); dst.layout = lease.layout(); dst.host_lease = std::move(lease); dst.pending = true; @@ -1549,19 +1617,9 @@ bool MoeHybridStreamEngine::activate_device_slot(int device_slot, if (err) *err = "SSD device slot is out of range"; return false; } - ScopedGpuDevice device_scope(runtime_->device); - if (!device_scope.ready()) { - if (err) *err = "failed to select SSD stream GPU"; - return false; - } Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) device_slot]; if (slot.pending) { - const cudaError_t gpu_err = cudaEventSynchronize(slot.ready); - if (gpu_err != cudaSuccess) { - if (err) *err = std::string("expert H2D synchronization failed: ") + - cudaGetErrorString(gpu_err); - return false; - } + ggml_backend_event_synchronize(slot.ready); slot.pending = false; slot.host_lease.reset(); } @@ -1620,11 +1678,6 @@ bool MoeHybridStreamEngine::warm_and_pin_device_cache( reserve_slots = std::max(2, reserve_slots); std::lock_guard compute_guard(runtime_->compute_mutex); - ScopedGpuDevice device_scope(runtime_->device); - if (!device_scope.ready()) { - if (err) *err = "failed to select streamed cache GPU for warmup"; - return false; - } std::vector candidates = entries; std::stable_sort(candidates.begin(), candidates.end(), @@ -2108,11 +2161,6 @@ bool eval_moe_streamed_experts( auto & runtime = *engine.runtime_; std::lock_guard compute_guard(runtime.compute_mutex); - ScopedGpuDevice device_scope(runtime.device); - if (!device_scope.ready()) { - if (err) *err = "failed to select streamed expert compute GPU"; - return false; - } if (!prepare_device_expert_layout( runtime, batch.layer, spec, err)) { return false; diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index d00ebba80..23a671e76 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -1196,7 +1196,7 @@ bool DeepSeek4Backend::init_hybrid_model() { stream_cold = info.cold_bytes > usable; if (stream_cold) nvme_device_cache_bytes = (size_t) usable; std::fprintf(stderr, - "[deepseek4] Strix cold tier: cold=%.2f GiB free=%.2f GiB " + "[deepseek4] secondary cold tier: cold=%.2f GiB free=%.2f GiB " "reserve=%.2f GiB warm-cache=%.2f GiB mode=%s\n", gib(info.cold_bytes), gib(expert_free), gib(reserve), gib(nvme_device_cache_bytes), diff --git a/server/src/kimi_k3/kimi_k3_backend.cpp b/server/src/kimi_k3/kimi_k3_backend.cpp index eb450a56f..04863edf1 100644 --- a/server/src/kimi_k3/kimi_k3_backend.cpp +++ b/server/src/kimi_k3/kimi_k3_backend.cpp @@ -1,5 +1,6 @@ #include "kimi_k3_backend.h" +#include "common/dynamic_backend.h" #include "common/moe_expert_package.h" #include "common/moe_hybrid_placement.h" #include "common/moe_stream_cache_policy.h" @@ -168,6 +169,7 @@ void KimiK3Backend::release_expert_backend() { expert_backend_ = nullptr; } expert_gpu_ = -1; + expert_backend_kind_ = PlacementBackend::Auto; } bool KimiK3Backend::init_streaming() { @@ -192,15 +194,39 @@ bool KimiK3Backend::init_streaming() { return false; } expert_gpu_ = owner.expert_gpu; - if (owner.heterogeneous()) { - expert_backend_ = ggml_backend_cuda_init(expert_gpu_); + const PlacementBackend primary_kind = + cfg_.device.backend == PlacementBackend::Auto + ? compiled_placement_backend() : cfg_.device.backend; + PlacementBackend expert_kind = primary_kind; + if (const char * raw = std::getenv("DFLASH_MOE_TP_BACKEND")) { + if (*raw && (!parse_placement_backend(raw, expert_kind) || + expert_kind == PlacementBackend::Auto)) { + std::fprintf(stderr, + "[kimi-k3] invalid DFLASH_MOE_TP_BACKEND=%s; " + "expected cuda or hip\n", raw); + return false; + } + } + const bool heterogeneous = + expert_kind != primary_kind || owner.heterogeneous(); + if (heterogeneous) { + expert_backend_ = init_placement_backend( + expert_kind, expert_gpu_, &error); if (!expert_backend_) { std::fprintf(stderr, - "[kimi-k3] expert backend init failed for device %d\n", - expert_gpu_); + "[kimi-k3] expert backend init failed for %s:%d: %s\n", + placement_backend_name(expert_kind), expert_gpu_, + error.c_str()); expert_gpu_ = -1; return false; } + expert_backend_kind_ = expert_kind; + std::fprintf(stderr, + "[kimi-k3] in-process routed owners primary=%s:%d " + "secondary=%s:%d transfer=backend-staged\n", + placement_backend_name(primary_kind), + cfg_.device.primary_gpu(), + placement_backend_name(expert_kind), expert_gpu_); } auto fail_streaming = [&]() { dual_stream_executor_.destroy(); @@ -240,7 +266,9 @@ bool KimiK3Backend::init_streaming() { routed_pool_bytes += bytes_per_expert * static_cast(weights_.n_expert); } - auto stream_config_for = [&](int gpu, const char * owner_name) { + auto stream_config_for = [&](ggml_backend_t owner_backend, int gpu, + PlacementBackend backend_kind, + const char * owner_name) { MoeStreamConfig stream_config = MoeStreamConfig::from_env(); if (std::getenv("DFLASH_MOE_NVME_DEVICE_CACHE_MB")) { stream_config.device_cache_bytes = @@ -249,7 +277,10 @@ bool KimiK3Backend::init_streaming() { } size_t free_bytes = 0; size_t total_bytes = 0; - ggml_backend_cuda_get_device_memory(gpu, &free_bytes, &total_bytes); + if (ggml_backend_dev_t device = + ggml_backend_get_device(owner_backend)) { + ggml_backend_dev_memory(device, &free_bytes, &total_bytes); + } const size_t gib = 1024ULL * 1024ULL * 1024ULL; const size_t reserve = std::max(2 * gib, total_bytes / 20); stream_config.device_cache_bytes = @@ -257,9 +288,9 @@ bool KimiK3Backend::init_streaming() { ? std::min(free_bytes - reserve, routed_pool_bytes) : 0; std::fprintf(stderr, - "[kimi-k3] %s streamed cache: gpu=%d free=%.2f GiB " + "[kimi-k3] %s streamed cache: device=%s:%d free=%.2f GiB " "reserve=%.2f GiB pool=%.2f GiB cache=%.2f GiB\n", - owner_name, gpu, + owner_name, placement_backend_name(backend_kind), gpu, static_cast(free_bytes) / gib, static_cast(reserve) / gib, static_cast(routed_pool_bytes) / gib, @@ -432,7 +463,7 @@ bool KimiK3Backend::init_streaming() { } const MoeStreamConfig primary_config = stream_config_for( - cfg_.device.primary_gpu(), "primary"); + backend_, cfg_.device.primary_gpu(), primary_kind, "primary"); if (!stream_engine_.init( backend_, max_streamed_expert_bytes, primary_config, &error)) { @@ -443,7 +474,7 @@ bool KimiK3Backend::init_streaming() { } if (expert_backend_) { const MoeStreamConfig secondary_config = stream_config_for( - expert_gpu_, "secondary"); + expert_backend_, expert_gpu_, expert_backend_kind_, "secondary"); if (!secondary_stream_engine_.init( expert_backend_, max_streamed_expert_bytes, secondary_config, &error)) { @@ -622,14 +653,16 @@ bool KimiK3Backend::init_streaming() { if (expert_backend_) { std::fprintf(stderr, "[kimi-k3] routed experts dual-owner: shards=%zu layers=%zu " - "primary=%d/%s/%.2fGiB secondary=%d/%s/%.2fGiB " + "primary=%s:%d/%s/%.2fGiB secondary=%s:%d/%s/%.2fGiB " "primary_share=%d/1000 placement=%s\n", weights_.shard_paths.size(), weights_.streamed_layer_regions.size(), - cfg_.device.primary_gpu(), stream_engine_.io_backend_name(), + placement_backend_name(primary_kind), cfg_.device.primary_gpu(), + stream_engine_.io_backend_name(), static_cast(stream_engine_.device_cache_bytes()) / (1024.0 * 1024.0 * 1024.0), - expert_gpu_, secondary_stream_engine_.io_backend_name(), + placement_backend_name(expert_backend_kind_), expert_gpu_, + secondary_stream_engine_.io_backend_name(), static_cast(secondary_stream_engine_.device_cache_bytes()) / (1024.0 * 1024.0 * 1024.0), stream_owner_policy_.primary_share_per_mille, diff --git a/server/src/kimi_k3/kimi_k3_backend.h b/server/src/kimi_k3/kimi_k3_backend.h index ac01d0065..f5472e992 100644 --- a/server/src/kimi_k3/kimi_k3_backend.h +++ b/server/src/kimi_k3/kimi_k3_backend.h @@ -17,10 +17,11 @@ struct KimiK3BackendConfig { const char * model_path = nullptr; DevicePlacement device; int stream_fd = -1; - // -1 resolves DFLASH_MOE_TP_GPU and otherwise keeps experts on the - // primary GPU. A different GPU becomes the secondary capacity owner; - // routed work is partitioned between both GPUs while dense KDA/MLA, - // recurrent state, and sampling remain primary-owned. + // -1 resolves DFLASH_MOE_TP_GPU and otherwise keeps the primary device + // index. DFLASH_MOE_TP_BACKEND may select a different in-process runtime + // (for example CUDA beside a HIP primary). A different backend or device + // becomes the secondary capacity owner; routed work is partitioned while + // dense KDA/MLA, recurrent state, and sampling remain primary-owned. int expert_gpu = -1; // Auto uses Kimi's capacity-safe file-backed routed experts. Resident is // retained as a deterministic oracle for small architecture fixtures. @@ -67,6 +68,7 @@ class KimiK3Backend final : public ModelBackend { KimiK3BackendConfig cfg_; ggml_backend_t backend_ = nullptr; ggml_backend_t expert_backend_ = nullptr; + PlacementBackend expert_backend_kind_ = PlacementBackend::Auto; int expert_gpu_ = -1; KimiK3Weights weights_; KimiK3Cache cache_; diff --git a/server/test/test_moe_stream_compute.cpp b/server/test/test_moe_stream_compute.cpp index cb7185988..63ac25487 100644 --- a/server/test/test_moe_stream_compute.cpp +++ b/server/test/test_moe_stream_compute.cpp @@ -1,7 +1,7 @@ #include "CppUnitTestFramework.hpp" +#include "common/dynamic_backend.h" #include "common/moe_hybrid_stream.h" -#include "ggml-cuda.h" #include "ggml-quants.h" #include @@ -648,13 +648,22 @@ TEST_CASE(MoeStreamComputeFixture, persistent_graph_matches_cpu_and_padded_mxfp4 if (const char * value = std::getenv("DFLASH_TEST_GPU")) { device = std::max(0, std::atoi(value)); } - if (device >= ggml_backend_cuda_get_device_count()) { - std::fprintf(stderr, "skip: requested CUDA/HIP device is unavailable\n"); - return; + PlacementBackend backend_kind = compiled_placement_backend(); + const char * backend_value = std::getenv("DFLASH_TEST_BACKEND"); + if (backend_value && *backend_value) { + STREAM_REQUIRE(parse_placement_backend(backend_value, backend_kind)); + STREAM_REQUIRE(backend_kind != PlacementBackend::Auto); } - ggml_backend_t backend = ggml_backend_cuda_init(device); + std::string init_error; + ggml_backend_t backend = init_placement_backend( + backend_kind, device, &init_error); if (!backend) { - std::fprintf(stderr, "skip: no CUDA/HIP backend available\n"); + if (backend_value && *backend_value) { + throw std::runtime_error( + "explicit stream test backend failed: " + init_error); + } + std::fprintf(stderr, "skip: no CUDA/HIP backend available: %s\n", + init_error.c_str()); return; } run_layout_case(backend, false); From d9ac9ca67bad6ab85bcf214a03a5c3139ef8f847 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 3 Aug 2026 23:35:12 +0200 Subject: [PATCH 18/20] test(kimi): benchmark backend-qualified heterogeneous owners --- server/docs/KIMI_K3_HETERO.md | 12 +-- .../scripts/benchmark_kimi_k3_deployments.py | 101 +++++++++++------- .../test_benchmark_kimi_k3_deployments.py | 46 +++++--- 3 files changed, 103 insertions(+), 56 deletions(-) diff --git a/server/docs/KIMI_K3_HETERO.md b/server/docs/KIMI_K3_HETERO.md index ee7912ba2..3bdfebb65 100644 --- a/server/docs/KIMI_K3_HETERO.md +++ b/server/docs/KIMI_K3_HETERO.md @@ -238,16 +238,16 @@ serially so they cannot contend for the SSD or GPUs: ```bash python3 server/scripts/benchmark_kimi_k3_deployments.py \ /models/Kimi-K3-UD-IQ1_S/Kimi-K3-UD-IQ1_S-00001-of-00014.gguf \ - --server-bin server/build-hip-dual/dflash_server \ - --r9700-device 0 --strix-device 1 \ + --server-bin server/build-hip-mixed/dflash_server \ + --primary-device hip:0 --secondary-device cuda:0 \ --output-dir bench-out/kimi-k3-lucebox ``` -The profiles are Strix-only+SSD and heterogeneous Strix+R9700+SSD. For the +The profiles are primary-only+SSD and heterogeneous primary+peer+SSD. For the published IQ1_S checkpoint, Strix is the capacity-safe primary because the -non-routed tensors alone exceed the R9700's memory; R9700 concurrently owns a -partition of routed experts. `--hetero-primary r9700` is available for a later -checkpoint whose non-routed plan fits the discrete GPU. +non-routed tensors exceed a 24/32 GiB discrete GPU; the RTX 3090 or R9700 can +concurrently own a partition of routed experts. Backend-qualified endpoints +also distinguish `hip:0` from `cuda:0` in the same process. Each profile starts from a fresh server, forces `--moe-storage ssd`, clears inherited MoE tuning, disables HTTP/prefix caches, and issues one cold followed diff --git a/server/scripts/benchmark_kimi_k3_deployments.py b/server/scripts/benchmark_kimi_k3_deployments.py index 974124187..c41f06ca4 100644 --- a/server/scripts/benchmark_kimi_k3_deployments.py +++ b/server/scripts/benchmark_kimi_k3_deployments.py @@ -28,6 +28,7 @@ _SPLIT_GGUF = re.compile(r"^(?P.+)-(?P\d+)-of-(?P\d+)(?P\.gguf)$") +_DEVICE_ENDPOINT = re.compile(r"^(?Phip|cuda):(?P\d+)$") _NVME_KEY_VALUE = re.compile( r"(?[a-z][a-z0-9-]*)=(?P\S+)" ) @@ -44,37 +45,54 @@ } +@dataclass(frozen=True) +class DeviceEndpoint: + backend: str + index: int + + def __post_init__(self) -> None: + if self.backend not in {"hip", "cuda"}: + raise ValueError(f"unsupported GPU backend: {self.backend}") + if self.index < 0: + raise ValueError("GPU device index must be non-negative") + + def __str__(self) -> str: + return f"{self.backend}:{self.index}" + + @dataclass(frozen=True) class Deployment: name: str - primary_device: int - secondary_device: int | None + primary: DeviceEndpoint + secondary: DeviceEndpoint | None + + +def parse_device_endpoint(value: str) -> DeviceEndpoint: + match = _DEVICE_ENDPOINT.fullmatch(value.strip().lower()) + if match is None: + raise argparse.ArgumentTypeError( + f"invalid GPU endpoint {value!r}; expected hip:N or cuda:N" + ) + return DeviceEndpoint(match.group("backend"), int(match.group("index"))) def build_deployments( profiles: list[str], - strix_device: int, - r9700_device: int, - hetero_primary: str, + primary: DeviceEndpoint, + secondary: DeviceEndpoint | None, ) -> list[Deployment]: - if strix_device < 0 or r9700_device < 0: - raise ValueError("GPU device indices must be non-negative") - if strix_device == r9700_device and "heterogeneous" in profiles: - raise ValueError("heterogeneous mode requires distinct R9700 and Strix devices") + if "heterogeneous" in profiles and secondary is None: + raise ValueError("heterogeneous mode requires --secondary-device") + if primary == secondary and "heterogeneous" in profiles: + raise ValueError("heterogeneous mode requires distinct GPU endpoints") deployments: list[Deployment] = [] for profile in profiles: - if profile == "strix-only": - deployments.append(Deployment("strix-only-ssd", strix_device, None)) + if profile in {"primary-only", "strix-only"}: + deployments.append(Deployment("primary-only-ssd", primary, None)) elif profile == "heterogeneous": - if hetero_primary == "strix": - deployments.append( - Deployment("heterogeneous-ssd", strix_device, r9700_device) - ) - else: - deployments.append( - Deployment("heterogeneous-ssd", r9700_device, strix_device) - ) + assert secondary is not None + deployments.append(Deployment("heterogeneous-ssd", primary, secondary)) else: raise ValueError(f"unknown deployment profile: {profile}") return deployments @@ -95,6 +113,7 @@ def deployment_environment( if key.startswith("DFLASH_MOE_NVME_") or key in { "DFLASH_MOE_STORAGE", "DFLASH_MOE_TP_GPU", + "DFLASH_MOE_TP_BACKEND", "DFLASH_MOE_PLACEMENT", "DFLASH_MOE_PRIMARY_SHARE_PER_MILLE", "DFLASH_MOE_DUAL_STREAM_TRACE", @@ -104,8 +123,9 @@ def deployment_environment( env["DFLASH_MOE_NVME_BACKEND"] = nvme_backend if device_cache_mb is not None: env["DFLASH_MOE_NVME_DEVICE_CACHE_MB"] = str(device_cache_mb) - if deployment.secondary_device is not None: - env["DFLASH_MOE_TP_GPU"] = str(deployment.secondary_device) + if deployment.secondary is not None: + env["DFLASH_MOE_TP_BACKEND"] = deployment.secondary.backend + env["DFLASH_MOE_TP_GPU"] = str(deployment.secondary.index) env["DFLASH_MOE_PRIMARY_SHARE_PER_MILLE"] = str(primary_share_per_mille) if placement is not None: env["DFLASH_MOE_PLACEMENT"] = str(placement) @@ -130,7 +150,7 @@ def server_command( "--port", str(port), "--target-device", - f"hip:{deployment.primary_device}", + str(deployment.primary), "--max-ctx", str(max_ctx), "--moe-storage", @@ -387,8 +407,8 @@ def run_deployment( "requests": [], } print( - f"\n[{deployment.name}] primary=hip:{deployment.primary_device} " - f"secondary={deployment.secondary_device}", + f"\n[{deployment.name}] primary={deployment.primary} " + f"secondary={deployment.secondary}", flush=True, ) process: subprocess.Popen[bytes] | None = None @@ -439,27 +459,33 @@ def run_deployment( def create_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( - description="Compare Kimi K3 Strix-only and heterogeneous SSD deployments" + description="Compare Kimi K3 primary-only and heterogeneous SSD deployments" ) parser.add_argument("model", type=Path, help="first shard of the split Kimi K3 GGUF") parser.add_argument( - "--server-bin", type=Path, default=Path("server/build-hip-dual/dflash_server") + "--server-bin", type=Path, default=Path("server/build-hip-mixed/dflash_server") ) parser.add_argument( "--profiles", nargs="+", - choices=("strix-only", "heterogeneous"), - default=("strix-only", "heterogeneous"), + choices=("primary-only", "heterogeneous", "strix-only"), + default=("primary-only", "heterogeneous"), + ) + parser.add_argument( + "--primary-device", + type=parse_device_endpoint, + required=True, + help=( + "backend-qualified device for dense/non-routed execution and the " + "primary expert partition" + ), ) - parser.add_argument("--strix-device", type=int, default=1) - parser.add_argument("--r9700-device", type=int, default=0) parser.add_argument( - "--hetero-primary", - choices=("strix", "r9700"), - default="strix", + "--secondary-device", + type=parse_device_endpoint, help=( - "primary device for the heterogeneous run; current IQ1_S defaults to Strix " - "because its non-routed tensors exceed R9700 VRAM" + "backend-qualified routed-expert capacity owner; mixed builds accept " + "a different runtime such as cuda:0" ), ) parser.add_argument("--primary-share-per-mille", type=int, default=500) @@ -509,9 +535,8 @@ def main() -> int: validate_args(args) deployments = build_deployments( list(args.profiles), - args.strix_device, - args.r9700_device, - args.hetero_primary, + args.primary_device, + args.secondary_device, ) if args.dry_run: resolved = [] diff --git a/server/scripts/test_benchmark_kimi_k3_deployments.py b/server/scripts/test_benchmark_kimi_k3_deployments.py index 2cf278d73..e554fc980 100644 --- a/server/scripts/test_benchmark_kimi_k3_deployments.py +++ b/server/scripts/test_benchmark_kimi_k3_deployments.py @@ -3,31 +3,41 @@ from pathlib import Path from benchmark_kimi_k3_deployments import ( + DeviceEndpoint, Deployment, build_deployments, deployment_environment, discover_model_files, extract_nvme_telemetry, + parse_device_endpoint, server_command, ) class KimiDeploymentBenchmarkTests(unittest.TestCase): def test_builds_capacity_safe_default_profiles(self): - profiles = build_deployments(["strix-only", "heterogeneous"], 1, 0, "strix") + strix = DeviceEndpoint("hip", 0) + rtx = DeviceEndpoint("cuda", 0) + profiles = build_deployments(["primary-only", "heterogeneous"], strix, rtx) self.assertEqual( profiles, [ - Deployment("strix-only-ssd", 1, None), - Deployment("heterogeneous-ssd", 1, 0), + Deployment("primary-only-ssd", strix, None), + Deployment("heterogeneous-ssd", strix, rtx), ], ) - def test_r9700_primary_is_explicit(self): - profiles = build_deployments(["heterogeneous"], 1, 0, "r9700") - self.assertEqual(profiles, [Deployment("heterogeneous-ssd", 0, 1)]) + def test_backend_qualified_endpoints_allow_matching_indices(self): + strix = parse_device_endpoint("hip:0") + rtx = parse_device_endpoint("CUDA:0") + profiles = build_deployments(["heterogeneous"], strix, rtx) + self.assertEqual( + profiles, [Deployment("heterogeneous-ssd", strix, rtx)] + ) with self.assertRaisesRegex(ValueError, "distinct"): - build_deployments(["heterogeneous"], 0, 0, "strix") + build_deployments(["heterogeneous"], strix, strix) + with self.assertRaisesRegex(ValueError, "secondary-device"): + build_deployments(["heterogeneous"], strix, None) def test_environment_removes_stale_tuning(self): base = { @@ -35,11 +45,16 @@ def test_environment_removes_stale_tuning(self): "DFLASH_MOE_STORAGE": "resident", "DFLASH_MOE_NVME_SLOTS": "64", "DFLASH_MOE_TP_GPU": "9", + "DFLASH_MOE_TP_BACKEND": "hip", "DFLASH_MOE_PLACEMENT": "/stale.json", } env = deployment_environment( base, - Deployment("heterogeneous-ssd", 1, 0), + Deployment( + "heterogeneous-ssd", + DeviceEndpoint("hip", 0), + DeviceEndpoint("cuda", 0), + ), "uring", 600, Path("/new.json"), @@ -51,6 +66,7 @@ def test_environment_removes_stale_tuning(self): self.assertNotIn("DFLASH_MOE_NVME_SLOTS", env) self.assertEqual(env["DFLASH_MOE_NVME_BACKEND"], "uring") self.assertEqual(env["DFLASH_MOE_NVME_DEVICE_CACHE_MB"], "4096") + self.assertEqual(env["DFLASH_MOE_TP_BACKEND"], "cuda") self.assertEqual(env["DFLASH_MOE_TP_GPU"], "0") self.assertEqual(env["DFLASH_MOE_PRIMARY_SHARE_PER_MILLE"], "600") self.assertEqual(env["DFLASH_MOE_PLACEMENT"], "/new.json") @@ -58,8 +74,8 @@ def test_environment_removes_stale_tuning(self): def test_strix_only_has_no_secondary_owner(self): env = deployment_environment( - {"DFLASH_MOE_TP_GPU": "0"}, - Deployment("strix-only-ssd", 1, None), + {"DFLASH_MOE_TP_GPU": "0", "DFLASH_MOE_TP_BACKEND": "cuda"}, + Deployment("primary-only-ssd", DeviceEndpoint("hip", 0), None), "auto", 500, None, @@ -67,10 +83,16 @@ def test_strix_only_has_no_secondary_owner(self): False, ) self.assertNotIn("DFLASH_MOE_TP_GPU", env) + self.assertNotIn("DFLASH_MOE_TP_BACKEND", env) command = server_command( - Path("server"), Path("model.gguf"), Deployment("strix", 1, None), 8080, 8192, [] + Path("server"), + Path("model.gguf"), + Deployment("strix", DeviceEndpoint("hip", 0), None), + 8080, + 8192, + [], ) - self.assertIn("hip:1", command) + self.assertIn("hip:0", command) self.assertEqual(command[command.index("--moe-storage") + 1], "ssd") self.assertEqual(command[command.index("--prefix-cache-slots") + 1], "0") From 2404599036912fedbb8e6ef49145ef43a070028e Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 4 Aug 2026 03:13:15 +0200 Subject: [PATCH 19/20] perf(moe): avoid UMA cache stalls and overlap cache hits --- .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 11 +- .../scripts/benchmark_kimi_k3_deployments.py | 28 ++++ .../test_benchmark_kimi_k3_deployments.py | 9 +- server/src/common/dynamic_backend.cpp | 25 +++ server/src/common/dynamic_backend.h | 10 ++ server/src/common/moe_hybrid_stream.cpp | 155 ++++++++++++++++-- server/src/common/moe_hybrid_stream.h | 5 + server/src/kimi_k3/kimi_k3_backend.cpp | 30 +++- server/test/test_mixed_cuda_hip.cpp | 16 ++ server/test/test_moe_stream_compute.cpp | 10 ++ 10 files changed, 275 insertions(+), 24 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 9b466fc3a..991e3bf27 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5382,10 +5382,14 @@ static bool ggml_backend_cuda_get_available_uma_memory(long * available_memory_k } #endif // defined(__linux__) +static void ggml_backend_cuda_device_get_native_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + ggml_backend_cuda_get_device_memory(ctx->device, free, total); +} + static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemGetInfo(free, total)); + ggml_backend_cuda_device_get_native_memory(dev, free, total); // ref: https://github.com/ggml-org/llama.cpp/pull/17368 #if defined(__linux__) @@ -6123,6 +6127,9 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_cuda_get_features; } + if (strcmp(name, "ggml_backend_dev_get_native_memory") == 0) { + return (void *)ggml_backend_cuda_device_get_native_memory; + } return nullptr; } diff --git a/server/scripts/benchmark_kimi_k3_deployments.py b/server/scripts/benchmark_kimi_k3_deployments.py index c41f06ca4..aad5146c2 100644 --- a/server/scripts/benchmark_kimi_k3_deployments.py +++ b/server/scripts/benchmark_kimi_k3_deployments.py @@ -106,6 +106,9 @@ def deployment_environment( placement: Path | None, device_cache_mb: int | None, dual_trace: bool, + nvme_direct: str = "auto", + nvme_slots: int | None = None, + nvme_cache_first: str = "on", ) -> dict[str, str]: """Return a clean MoE environment without inherited benchmark tuning.""" env = dict(base) @@ -121,6 +124,12 @@ def deployment_environment( env.pop(key) env["DFLASH_MOE_NVME_BACKEND"] = nvme_backend + env["DFLASH_MOE_NVME_DIRECT"] = nvme_direct + env["DFLASH_MOE_NVME_CACHE_FIRST"] = ( + "1" if nvme_cache_first == "on" else "0" + ) + if nvme_slots is not None: + env["DFLASH_MOE_NVME_SLOTS"] = str(nvme_slots) if device_cache_mb is not None: env["DFLASH_MOE_NVME_DEVICE_CACHE_MB"] = str(device_cache_mb) if deployment.secondary is not None: @@ -389,6 +398,9 @@ def run_deployment( args.placement, args.device_cache_mb, args.dual_trace, + args.nvme_direct, + args.nvme_slots, + args.nvme_cache_first, ) command = server_command( args.server_bin, @@ -494,6 +506,17 @@ def create_parser() -> argparse.ArgumentParser: parser.add_argument( "--nvme-backend", choices=("auto", "uring", "pread", "mmap"), default="auto" ) + parser.add_argument( + "--nvme-direct", choices=("auto", "on", "off"), default="auto" + ) + parser.add_argument( + "--nvme-slots", + type=int, + help="page-locked host slots per SSD expert owner (engine default: 8)", + ) + parser.add_argument( + "--nvme-cache-first", choices=("on", "off"), default="on" + ) parser.add_argument("--port", type=int, default=18080) parser.add_argument("--max-ctx", type=int, default=8192) parser.add_argument("--max-tokens", type=int, default=8) @@ -522,6 +545,8 @@ def validate_args(args: argparse.Namespace) -> None: raise ValueError("primary-share-per-mille must be in [0, 1000]") if args.device_cache_mb is not None and args.device_cache_mb < 0: raise ValueError("device-cache-mb must be non-negative") + if args.nvme_slots is not None and not 2 <= args.nvme_slots <= 64: + raise ValueError("nvme-slots must be in [2, 64]") if len(set(args.profiles)) != len(args.profiles): raise ValueError("deployment profiles must not be repeated") if args.placement is not None and not args.dry_run and not args.placement.is_file(): @@ -549,6 +574,9 @@ def main() -> int: args.placement, args.device_cache_mb, args.dual_trace, + args.nvme_direct, + args.nvme_slots, + args.nvme_cache_first, ) resolved.append( { diff --git a/server/scripts/test_benchmark_kimi_k3_deployments.py b/server/scripts/test_benchmark_kimi_k3_deployments.py index e554fc980..be8a14611 100644 --- a/server/scripts/test_benchmark_kimi_k3_deployments.py +++ b/server/scripts/test_benchmark_kimi_k3_deployments.py @@ -60,11 +60,16 @@ def test_environment_removes_stale_tuning(self): Path("/new.json"), 4096, True, + "off", + 16, + "off", ) self.assertEqual(env["PATH"], "/bin") self.assertNotIn("DFLASH_MOE_STORAGE", env) - self.assertNotIn("DFLASH_MOE_NVME_SLOTS", env) + self.assertEqual(env["DFLASH_MOE_NVME_SLOTS"], "16") self.assertEqual(env["DFLASH_MOE_NVME_BACKEND"], "uring") + self.assertEqual(env["DFLASH_MOE_NVME_DIRECT"], "off") + self.assertEqual(env["DFLASH_MOE_NVME_CACHE_FIRST"], "0") self.assertEqual(env["DFLASH_MOE_NVME_DEVICE_CACHE_MB"], "4096") self.assertEqual(env["DFLASH_MOE_TP_BACKEND"], "cuda") self.assertEqual(env["DFLASH_MOE_TP_GPU"], "0") @@ -84,6 +89,8 @@ def test_strix_only_has_no_secondary_owner(self): ) self.assertNotIn("DFLASH_MOE_TP_GPU", env) self.assertNotIn("DFLASH_MOE_TP_BACKEND", env) + self.assertEqual(env["DFLASH_MOE_NVME_DIRECT"], "auto") + self.assertEqual(env["DFLASH_MOE_NVME_CACHE_FIRST"], "1") command = server_command( Path("server"), Path("model.gguf"), diff --git a/server/src/common/dynamic_backend.cpp b/server/src/common/dynamic_backend.cpp index 634a6ce4c..60c894f83 100644 --- a/server/src/common/dynamic_backend.cpp +++ b/server/src/common/dynamic_backend.cpp @@ -175,6 +175,31 @@ PlacementBackend placement_backend_of(ggml_backend_t backend) { return PlacementBackend::Auto; } +bool backend_native_memory(ggml_backend_t backend, + size_t * free_bytes, + size_t * total_bytes) { + if (!backend || !free_bytes || !total_bytes) return false; + ggml_backend_dev_t device = ggml_backend_get_device(backend); + if (!device) return false; + ggml_backend_reg_t registry = ggml_backend_dev_backend_reg(device); + if (!registry) return false; + + using query_fn = void (*)(ggml_backend_dev_t, size_t *, size_t *); + auto query = reinterpret_cast(ggml_backend_reg_get_proc_address( + registry, "ggml_backend_dev_get_native_memory")); + if (!query) return false; + + size_t native_free = 0; + size_t native_total = 0; + query(device, &native_free, &native_total); + if (native_free == 0 || native_total == 0 || native_free > native_total) { + return false; + } + *free_bytes = native_free; + *total_bytes = native_total; + return true; +} + BackendPairCapabilities backend_pair_capabilities(ggml_backend_t first, ggml_backend_t second) { BackendPairCapabilities result; diff --git a/server/src/common/dynamic_backend.h b/server/src/common/dynamic_backend.h index 77a6186b9..089d20229 100644 --- a/server/src/common/dynamic_backend.h +++ b/server/src/common/dynamic_backend.h @@ -7,6 +7,7 @@ #include "ggml-backend.h" +#include #include namespace dflash::common { @@ -29,6 +30,15 @@ ggml_backend_t init_placement_backend(PlacementBackend backend, // Resolve the vendor represented by an initialized ggml backend. PlacementBackend placement_backend_of(ggml_backend_t backend); +// Return the largest-memory-pool view reported by the runtime itself. This is +// distinct from ggml_backend_dev_memory(): on unified-memory GPUs that generic +// query may deliberately report available system RAM, while a single native +// device allocation is still bounded by the runtime's contiguous pool. +// Returns false when a backend does not expose the optional query. +bool backend_native_memory(ggml_backend_t backend, + size_t * free_bytes, + size_t * total_bytes); + // Describe the operations that are safe between two initialized backends. // This deliberately keys off backend identity rather than vendor names so the // scheduling code also works for future runtime modules. diff --git a/server/src/common/moe_hybrid_stream.cpp b/server/src/common/moe_hybrid_stream.cpp index 08e456e2a..cef2a036f 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -189,6 +189,9 @@ MoeStreamConfig MoeStreamConfig::from_env() { "DFLASH_MOE_NVME_GRAPH_CACHE", config.graph_cache_entries, 0, 64); config.fused_decode = env_bounded_int( "DFLASH_MOE_NVME_FUSED_DECODE", config.fused_decode ? 1 : 0, 0, 1) != 0; + config.cache_first_decode = env_bounded_int( + "DFLASH_MOE_NVME_CACHE_FIRST", + config.cache_first_decode ? 1 : 0, 0, 1) != 0; config.device_cache_bytes = env_mib( "DFLASH_MOE_NVME_DEVICE_CACHE_MB", config.device_cache_bytes); config.prefill_threshold = env_bounded_int( @@ -867,6 +870,12 @@ struct MoeHybridStreamEngine::Runtime { size_t max_expert_bytes = 0; MoeStreamConfig config{}; BackendHostAllocator host_allocator; + // Quantized CUDA/HIP kernels may read backend-added row padding. Keep one + // immutable backend-pinned zero source and upload only those padding bytes + // with each expert. Clearing a capacity-sized cache eagerly can fault tens + // of GiB of APU managed memory before the first request. + void * zero_padding = nullptr; + size_t zero_padding_bytes = 0; std::unique_ptr io; ggml_backend_buffer_t device_pool_buffer = nullptr; void * device_pool = nullptr; @@ -921,6 +930,37 @@ void release_device_cache(RuntimeT & runtime) { runtime.device_pool_bytes = 0; } +template +bool ensure_zero_padding(RuntimeT & runtime, size_t required_bytes, + std::string * err) { + if (required_bytes == 0 || required_bytes <= runtime.zero_padding_bytes) { + return true; + } + const size_t allocation_bytes = align_up(required_bytes, 256); + if (allocation_bytes == 0) { + if (err) *err = "SSD quantization padding size overflow"; + return false; + } + + void * replacement = nullptr; + if (!runtime.host_allocator.allocate(&replacement, allocation_bytes)) { + if (err) *err = "failed to allocate SSD quantization padding staging"; + return false; + } + std::memset(replacement, 0, allocation_bytes); + + // The previous immutable source may still be referenced by queued H2D + // copies. Growth normally happens once, on the first numerical layout, so + // synchronize only before replacement and never on the steady-state path. + if (runtime.zero_padding) { + ggml_backend_synchronize(runtime.transfer_backend); + runtime.host_allocator.release(runtime.zero_padding); + } + runtime.zero_padding = replacement; + runtime.zero_padding_bytes = allocation_bytes; + return true; +} + template bool allocate_device_cache(RuntimeT & runtime, std::string * err, size_t minimum_stride = 0) { @@ -997,11 +1037,6 @@ bool allocate_device_cache(RuntimeT & runtime, std::string * err, release_device_cache(runtime); return false; } - ggml_backend_buffer_clear(runtime.device_pool_buffer, 0); - // The pool is shared by the compute and upload backend instances. Finish - // initialization on the compute stream before the upload stream can reuse - // any slot, otherwise zero-filled quantization padding could race H2D. - ggml_backend_synchronize(runtime.backend); for (size_t i = 0; i < runtime.device_slots.size(); ++i) { auto & slot = runtime.device_slots[i]; slot.data = base + i * runtime.device_stride; @@ -1318,7 +1353,8 @@ void MoeHybridStreamEngine::destroy() { "device-cache=%.1f MiB slots=%zu pinned=%zu " "hits=%llu misses=%llu evictions=%llu " "graphs=%llu graph-hits=%llu graph-evictions=%llu launches=%llu " - "fused-decode-launches=%llu fused-decode-experts=%llu\n", + "fused-decode-launches=%llu fused-decode-experts=%llu " + "cache-first-reorders=%llu cache-first-experts=%llu\n", runtime_->io->effective_backend_name(), (unsigned long long) stats.requests, (unsigned long long) stats.read_ops, @@ -1339,11 +1375,18 @@ void MoeHybridStreamEngine::destroy() { (unsigned long long) runtime_->compute_stats.graph_evictions, (unsigned long long) runtime_->compute_stats.graph_launches, (unsigned long long) runtime_->compute_stats.fused_decode_launches, - (unsigned long long) runtime_->compute_stats.fused_decode_experts); + (unsigned long long) runtime_->compute_stats.fused_decode_experts, + (unsigned long long) runtime_->compute_stats.cache_first_reorders, + (unsigned long long) runtime_->compute_stats.cache_first_experts); } runtime_->io->destroy(); } runtime_->io.reset(); + if (runtime_->zero_padding) { + runtime_->host_allocator.release(runtime_->zero_padding); + runtime_->zero_padding = nullptr; + runtime_->zero_padding_bytes = 0; + } if (runtime_->transfer_backend) { ggml_backend_free(runtime_->transfer_backend); runtime_->transfer_backend = nullptr; @@ -1472,6 +1515,16 @@ bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, if (err) *err = "prepared streamed expert exceeds GPU device stride"; return false; } + size_t maximum_padding = 0; + for (int i = 0; i < device_layout.component_count; ++i) { + const Runtime::DeviceComponentLayout & component = + device_layout.components[i]; + maximum_padding = std::max( + maximum_padding, + component.alloc_bytes - component.logical_bytes); + } + if (!ensure_zero_padding(*runtime_, maximum_padding, err)) return false; + for (int i = 0; i < device_layout.component_count; ++i) { const Runtime::DeviceComponentLayout & device_component = device_layout.components[i]; @@ -1506,6 +1559,15 @@ bool MoeHybridStreamEngine::stage_expert_async(int layer, int expert_id, runtime_->transfer_backend, dst.transfer_tensor, source, device_component.offset, device_component.logical_bytes); + const size_t padding_bytes = + device_component.alloc_bytes - device_component.logical_bytes; + if (padding_bytes > 0) { + ggml_backend_tensor_set_async( + runtime_->transfer_backend, dst.transfer_tensor, + runtime_->zero_padding, + device_component.offset + device_component.logical_bytes, + padding_bytes); + } } dst.device_layout = device_layout; } else { @@ -2355,9 +2417,54 @@ bool eval_moe_streamed_experts( return true; } + // The weighted expert sum is order-independent mathematically. Execute + // device-resident experts first so their compute overlaps every admitted + // SSD miss instead of blocking on the first cold expert encountered by ID. + // Contributions are accumulated later in the original deterministic order + // to preserve the previous floating-point result. + std::vector execution_experts = unique_experts; + bool cache_first_reordered = false; + if (runtime.config.cache_first_decode && batch.n_tokens == 1 && + execution_experts.size() > 1) { + const auto is_device_resident = [&](int32_t expert) { + const auto found = runtime.device_index.find( + device_key(batch.layer, expert)); + if (found == runtime.device_index.end() || found->second < 0 || + found->second >= (int) runtime.device_slots.size()) { + return false; + } + const auto & slot = + runtime.device_slots[(size_t) found->second]; + return slot.valid && slot.cache_managed && + slot.compute_users == 0 && slot.key.layer == batch.layer && + slot.key.expert == expert; + }; + const auto cold_begin = std::stable_partition( + execution_experts.begin(), execution_experts.end(), + is_device_resident); + const size_t resident_count = (size_t) std::distance( + execution_experts.begin(), cold_begin); + cache_first_reordered = execution_experts != unique_experts; + if (cache_first_reordered) { + ++runtime.compute_stats.cache_first_reorders; + runtime.compute_stats.cache_first_experts += resident_count; + } + } + + std::vector ordered_contributions; + if (cache_first_reordered) { + size_t contribution_values = 0; + if (!checked_mul_size(output_values, unique_experts.size(), + contribution_values)) { + if (err) *err = "streamed expert contribution size overflow"; + return false; + } + ordered_contributions.assign(contribution_values, 0.0f); + } + int staged_slot = -1; if (!engine.stage_expert_cached_async( - batch.layer, unique_experts[0], &staged_slot, err)) { + batch.layer, execution_experts[0], &staged_slot, err)) { return false; } @@ -2453,7 +2560,7 @@ bool eval_moe_streamed_experts( std::vector result; for (size_t expert_index = 0; - expert_index < unique_experts.size(); ++expert_index) { + expert_index < execution_experts.size(); ++expert_index) { const int current_slot = staged_slot; if (!engine.activate_device_slot(current_slot, err)) return false; auto release_current = [&]() { @@ -2461,7 +2568,7 @@ bool eval_moe_streamed_experts( }; hits.clear(); - const int32_t expert = unique_experts[expert_index]; + const int32_t expert = execution_experts[expert_index]; for (int token = 0; token < batch.n_tokens; ++token) { float combined_weight = 0.0f; for (int rank = 0; rank < batch.top_k; ++rank) { @@ -2518,10 +2625,10 @@ bool eval_moe_streamed_experts( // Compute N is running while the already-admitted read for N+1 is // acquired and uploaded into a different, eviction-protected slot. - if (expert_index + 1 < unique_experts.size()) { + if (expert_index + 1 < execution_experts.size()) { int next_slot = -1; if (!engine.stage_expert_cached_async( - batch.layer, unique_experts[expert_index + 1], + batch.layer, execution_experts[expert_index + 1], &next_slot, err)) { ggml_backend_synchronize(runtime.backend); release_current(); @@ -2534,9 +2641,19 @@ bool eval_moe_streamed_experts( release_current(); return false; } + const size_t original_expert_index = cache_first_reordered + ? (size_t) std::distance( + unique_experts.begin(), + std::lower_bound( + unique_experts.begin(), unique_experts.end(), expert)) + : 0; for (size_t i = 0; i < hits.size(); ++i) { - float * dst = out.data() + - (size_t) hits[i].token * (size_t) spec.output_dim; + float * dst = cache_first_reordered + ? ordered_contributions.data() + + original_expert_index * output_values + + (size_t) hits[i].token * (size_t) spec.output_dim + : out.data() + + (size_t) hits[i].token * (size_t) spec.output_dim; const float * src = result.data() + i * (size_t) spec.output_dim; const float weight = hits[i].weight; @@ -2546,6 +2663,16 @@ bool eval_moe_streamed_experts( } release_current(); } + if (cache_first_reordered) { + for (size_t expert_index = 0; + expert_index < unique_experts.size(); ++expert_index) { + const float * contribution = ordered_contributions.data() + + expert_index * output_values; + for (size_t value = 0; value < output_values; ++value) { + out[value] += contribution[value]; + } + } + } return true; } diff --git a/server/src/common/moe_hybrid_stream.h b/server/src/common/moe_hybrid_stream.h index f3e9b32c2..733728e97 100644 --- a/server/src/common/moe_hybrid_stream.h +++ b/server/src/common/moe_hybrid_stream.h @@ -34,6 +34,9 @@ struct MoeStreamConfig { // backend graph and reduce on the GPU, avoiding one synchronization and // D2H copy per expert. Misses and multi-token prefill retain the pipeline. bool fused_decode = true; + // On a partial device-cache hit, compute resident decode experts before + // waiting on admitted SSD reads. Contributions retain their old sum order. + bool cache_first_decode = true; // Optional adaptive GPU expert-cache budget. Zero keeps only the pipeline // slots. The hardware planner can safely assign otherwise-unused Strix // memory here while retaining its KV/graph reserve. @@ -107,6 +110,8 @@ struct MoeStreamComputeStats { uint64_t graph_launches = 0; uint64_t fused_decode_launches = 0; uint64_t fused_decode_experts = 0; + uint64_t cache_first_reorders = 0; + uint64_t cache_first_experts = 0; }; struct MoeStreamCacheWarmEntry { diff --git a/server/src/kimi_k3/kimi_k3_backend.cpp b/server/src/kimi_k3/kimi_k3_backend.cpp index 04863edf1..7684fb6ba 100644 --- a/server/src/kimi_k3/kimi_k3_backend.cpp +++ b/server/src/kimi_k3/kimi_k3_backend.cpp @@ -281,18 +281,34 @@ bool KimiK3Backend::init_streaming() { ggml_backend_get_device(owner_backend)) { ggml_backend_dev_memory(device, &free_bytes, &total_bytes); } + // ggml's generic UMA query intentionally reports available system RAM + // so model loaders can use managed memory. The streamed cache needs one + // contiguous native allocation, so cap it using the runtime's native + // pool even when the owner is provided by a dynamically loaded peer. + size_t allocation_free = free_bytes; + size_t allocation_total = total_bytes; + backend_native_memory( + owner_backend, &allocation_free, &allocation_total); const size_t gib = 1024ULL * 1024ULL * 1024ULL; - const size_t reserve = std::max(2 * gib, total_bytes / 20); - stream_config.device_cache_bytes = - free_bytes > reserve - ? std::min(free_bytes - reserve, routed_pool_bytes) - : 0; + const size_t system_reserve = + std::max(2 * gib, total_bytes / 20); + const size_t allocation_reserve = + std::max(2 * gib, allocation_total / 20); + const size_t system_budget = free_bytes > system_reserve + ? free_bytes - system_reserve : 0; + const size_t allocation_budget = allocation_free > allocation_reserve + ? allocation_free - allocation_reserve : 0; + stream_config.device_cache_bytes = std::min( + {system_budget, allocation_budget, routed_pool_bytes}); std::fprintf(stderr, "[kimi-k3] %s streamed cache: device=%s:%d free=%.2f GiB " - "reserve=%.2f GiB pool=%.2f GiB cache=%.2f GiB\n", + "alloc-free=%.2f GiB reserve=%.2f/%.2f GiB pool=%.2f GiB " + "cache=%.2f GiB\n", owner_name, placement_backend_name(backend_kind), gpu, static_cast(free_bytes) / gib, - static_cast(reserve) / gib, + static_cast(allocation_free) / gib, + static_cast(system_reserve) / gib, + static_cast(allocation_reserve) / gib, static_cast(routed_pool_bytes) / gib, static_cast(stream_config.device_cache_bytes) / gib); return stream_config; diff --git a/server/test/test_mixed_cuda_hip.cpp b/server/test/test_mixed_cuda_hip.cpp index 81b73e6f7..2b39c256a 100644 --- a/server/test/test_mixed_cuda_hip.cpp +++ b/server/test/test_mixed_cuda_hip.cpp @@ -12,12 +12,26 @@ #include using dflash::common::PlacementBackend; +using dflash::common::backend_native_memory; using dflash::common::backend_pair_capabilities; using dflash::common::init_placement_backend; using dflash::common::placement_backend_of; namespace { +bool run_native_memory_query(ggml_backend_t backend, const char * label) { + size_t free_bytes = 0; + size_t total_bytes = 0; + const bool ok = backend_native_memory( + backend, &free_bytes, &total_bytes) && + free_bytes > 0 && total_bytes >= free_bytes; + std::printf("mixed-backend %s native memory: %s (%.2f/%.2f GiB)\n", + label, ok ? "ok" : "FAILED", + (double) free_bytes / (1024.0 * 1024.0 * 1024.0), + (double) total_bytes / (1024.0 * 1024.0 * 1024.0)); + return ok; +} + bool run_scale(ggml_backend_t backend, const char * label) { constexpr int64_t n = 4096; ggml_init_params params{}; @@ -491,6 +505,8 @@ int main() { ok = placement_backend_of(cuda) == PlacementBackend::Cuda && placement_backend_of(hip) == PlacementBackend::Hip && !pair.same_runtime && !pair.native_gpu_handoff && ok; + ok = run_native_memory_query(cuda, "CUDA") && ok; + ok = run_native_memory_query(hip, "HIP") && ok; ok = run_scale(cuda, "CUDA") && ok; ok = run_scale(hip, "HIP") && ok; ok = run_cross_copy(cuda, hip, "CUDA->HIP") && ok; diff --git a/server/test/test_moe_stream_compute.cpp b/server/test/test_moe_stream_compute.cpp index 63ac25487..c4980a099 100644 --- a/server/test/test_moe_stream_compute.cpp +++ b/server/test/test_moe_stream_compute.cpp @@ -354,6 +354,14 @@ void run_fused_decode_case(ggml_backend_t backend, bool mxfp4) { &prepare_stats, &error)); STREAM_REQUIRE(prepare_stats.capacity_drops == 1); + // Make the largest selected ID resident. The cold fallback must execute it + // before lower-ID misses while preserving the original accumulation order. + int resident_slot = -1; + STREAM_REQUIRE(engine.stage_expert_cached_async( + 0, 2, &resident_slot, &error)); + STREAM_REQUIRE(engine.activate_device_slot(resident_slot, &error)); + engine.release_device_slot(resident_slot); + const std::vector expected = cpu_reference( gate_reference, up_reference, down_reference, input, ids, weights, 1, kDecodeTopK); @@ -375,6 +383,8 @@ void run_fused_decode_case(ggml_backend_t backend, bool mxfp4) { const MoeStreamComputeStats cold = engine.compute_stats(); STREAM_REQUIRE(cold.graph_launches == kDecodeTopK); STREAM_REQUIRE(cold.fused_decode_launches == 0); + STREAM_REQUIRE(cold.cache_first_reorders == 1); + STREAM_REQUIRE(cold.cache_first_experts == 1); // Populate every slot, then verify the all-resident fused path. for (int expert = 0; expert < kExperts; ++expert) { From f34eba6057b7d908615246dc99688236f582547d Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:26:46 +0200 Subject: [PATCH 20/20] fix(kimi): suppress Windows min max macros --- server/src/kimi_k3/kimi_k3_backend.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/server/src/kimi_k3/kimi_k3_backend.cpp b/server/src/kimi_k3/kimi_k3_backend.cpp index 7684fb6ba..c34bb7715 100644 --- a/server/src/kimi_k3/kimi_k3_backend.cpp +++ b/server/src/kimi_k3/kimi_k3_backend.cpp @@ -1,3 +1,7 @@ +#if defined(_WIN32) && !defined(NOMINMAX) +#define NOMINMAX +#endif + #include "kimi_k3_backend.h" #include "common/dynamic_backend.h"