diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3cfe3e446..28a0ca994 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -79,12 +79,15 @@ 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 \ + 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" --no-tests=error + ctest --output-on-failure \ + -R "server_unit|deepseek4_unit|test_moe_nvme_scheduler|test_moe_expert_package" \ + --no-tests=error - name: Populate venv with cu128 torch + setuptools # First pass: install the workspace's default deps. dflash declares @@ -174,6 +177,8 @@ jobs: -DCMAKE_BUILD_TYPE=Release 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 @@ -184,6 +189,14 @@ 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 + + - name: Run NVMe scheduler production tests + 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 # ~16 GB Qwen3.6-27B target + draft (~1-2 min). The weights are already @@ -290,11 +303,13 @@ 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_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|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 c295aed46..1401112e5 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -364,6 +364,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 ) @@ -396,6 +397,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 @@ -434,6 +439,9 @@ 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_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 @@ -757,6 +765,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) @@ -806,6 +818,107 @@ 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() + + 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. + 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() + + # 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 + ${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") + 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 + # 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 @@ -1056,6 +1169,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}) @@ -1235,6 +1359,8 @@ 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_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/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/docs/DS4.md b/server/docs/DS4.md index 3c1787421..3045771bc 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -225,6 +225,18 @@ 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 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. `--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. + ## Environment Variables | Variable | Purpose | @@ -245,6 +257,9 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. | `DFLASH_DS4_COMP_PAD_STRIDE` | Exact compressed-KV padding bucket; wider buckets trade small masked work for fewer verifier graph captures. | | `DFLASH_DS4_DISABLE_GROUPED_OUTPUT_PROJECTION` | Diagnostic fallback for runtimes that cannot preserve grouped projection metadata across a scheduler copy. | | `DFLASH_CUDA_BACKEND_PATH` / `DFLASH_HIP_BACKEND_PATH` | Optional explicit peer backend module path. | +| `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. | | `GGML_BATCH_PEER_COPIES` | Batch peer-runtime copies and unlike-runtime pinned-host staging with one source wait per split. The old `GGML_CUDA_BATCH_PEER_COPIES` spelling remains an alias. | diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index 77f308e1f..3709c025d 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -40,6 +40,15 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_CUDA_BACKEND_PATH` / `DFLASH_HIP_BACKEND_PATH` | auto-discovered beside the executable | Explicit peer module file path for a mixed CUDA+HIP build. | | `GGML_BATCH_PEER_COPIES` | unset | BURN-IN: batch peer-runtime copies and unlike-runtime host staging with one source wait per split. `GGML_CUDA_BATCH_PEER_COPIES` remains a compatibility alias. | | `GGML_SCHED_PROFILE` / `GGML_SCHED_PROFILE_MIN_SPLITS` | unset / 1 | DEBUG: report scheduler splits, copy volume, submission time, and source/destination synchronization time. | +| `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_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. | | `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. | | `DFLASH_MMID_TELEMETRY` | unset | DEBUG: report MUL_MAT_ID dispatch, MMVQ variant, and per-node graph compatibility. | @@ -181,6 +190,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 @@ -191,16 +201,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_MOE_TP_BACKEND` - deepseek4_backend.cpp - `DFLASH_NO_MASK` - laguna_backend.cpp - `DFLASH_NO_MOE_ROUTER_FUSE` - qwen35moe_ffn.cpp diff --git a/server/docs/KIMI_K3_HETERO.md b/server/docs/KIMI_K3_HETERO.md new file mode 100644 index 000000000..3bdfebb65 --- /dev/null +++ b/server/docs/KIMI_K3_HETERO.md @@ -0,0 +1,257 @@ +# 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 + +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 + 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. +- 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. +- Routed-expert ownership is independent of contiguous layer splitting. + 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, + `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 +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. + +The SSD text path is implemented, but two qualification boundaries remain: + +- 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. + +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 + +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, 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 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, +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. + +## 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 +R9700, or 156.94 GiB of unique physical weight capacity before runtime +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. 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. +Under a uniform balanced-routing assumption this covers about 14-17% of the +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 optimization milestone + +The implementation no longer needs another generic cache, a second Kimi +router, per-layer worker creation, or a first full-scale smoke test. The next +work is: + +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. +3. Compare end-to-end output and token rate against ordinary llama.cpp + CPU/GPU offload. +4. Replace the correctness-first host partial join with a device-resident peer + join, then tune the placement until the two owner branches balance. + +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-mixed/dflash_server \ + --primary-device hip:0 --secondary-device cuda:0 \ + --output-dir bench-out/kimi-k3-lucebox +``` + +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 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 +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/docs/MOE_NVME_STREAMING.md b/server/docs/MOE_NVME_STREAMING.md new file mode 100644 index 000000000..62320a784 --- /dev/null +++ b/server/docs/MOE_NVME_STREAMING.md @@ -0,0 +1,400 @@ +# 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. + +## Deployment shapes + +On a full Lucebox, the three owners have distinct jobs: + +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 + 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 +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. + +`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: + +- 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 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. +- 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. + +## 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 +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 + +./build-hip-dual/dflash_server /path/to/model.gguf \ + --target-device hip:0 --peer-access --moe-storage auto +``` + +`--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. + +### 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 + +./build-hip/dflash_server /path/to/model.gguf \ + --target-device hip:0 --max-ctx 8192 --moe-storage ssd +``` + +`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; +`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 --moe-storage auto +``` + +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. 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 + +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 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 +SSD slots/caches and execute their selected routed experts concurrently: + +```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`. `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 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 + +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_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 | +| `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 | +| `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 | +| `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 | + +Shutdown telemetry reports logical and physical bytes, measured read service +rate, cache hits, demand wait/timeouts, de-duplication, dropped speculation, +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 +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 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 +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 +environment/default owner. + +## 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. + +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. + +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. + +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 +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 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/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/scripts/benchmark_kimi_k3_deployments.py b/server/scripts/benchmark_kimi_k3_deployments.py new file mode 100644 index 000000000..aad5146c2 --- /dev/null +++ b/server/scripts/benchmark_kimi_k3_deployments.py @@ -0,0 +1,642 @@ +#!/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)$") +_DEVICE_ENDPOINT = re.compile(r"^(?Phip|cuda):(?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) +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: 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], + primary: DeviceEndpoint, + secondary: DeviceEndpoint | None, +) -> list[Deployment]: + 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 in {"primary-only", "strix-only"}: + deployments.append(Deployment("primary-only-ssd", primary, None)) + elif profile == "heterogeneous": + assert secondary is not None + deployments.append(Deployment("heterogeneous-ssd", primary, secondary)) + 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, + 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) + for key in list(env): + 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", + }: + 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: + 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) + 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", + str(deployment.primary), + "--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 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 + + +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, + args.nvme_direct, + args.nvme_slots, + args.nvme_cache_first, + ) + 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={deployment.primary} " + f"secondary={deployment.secondary}", + 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 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-mixed/dflash_server") + ) + parser.add_argument( + "--profiles", + nargs="+", + 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( + "--secondary-device", + type=parse_device_endpoint, + help=( + "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) + 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( + "--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) + 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 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(): + 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.primary_device, + args.secondary_device, + ) + 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, + args.nvme_direct, + args.nvme_slots, + args.nvme_cache_first, + ) + 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..be8a14611 --- /dev/null +++ b/server/scripts/test_benchmark_kimi_k3_deployments.py @@ -0,0 +1,155 @@ +import tempfile +import unittest +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): + strix = DeviceEndpoint("hip", 0) + rtx = DeviceEndpoint("cuda", 0) + profiles = build_deployments(["primary-only", "heterogeneous"], strix, rtx) + self.assertEqual( + profiles, + [ + Deployment("primary-only-ssd", strix, None), + Deployment("heterogeneous-ssd", strix, rtx), + ], + ) + + 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"], strix, strix) + with self.assertRaisesRegex(ValueError, "secondary-device"): + build_deployments(["heterogeneous"], strix, None) + + 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_TP_BACKEND": "hip", + "DFLASH_MOE_PLACEMENT": "/stale.json", + } + env = deployment_environment( + base, + Deployment( + "heterogeneous-ssd", + DeviceEndpoint("hip", 0), + DeviceEndpoint("cuda", 0), + ), + "uring", + 600, + Path("/new.json"), + 4096, + True, + "off", + 16, + "off", + ) + self.assertEqual(env["PATH"], "/bin") + self.assertNotIn("DFLASH_MOE_STORAGE", 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") + 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", "DFLASH_MOE_TP_BACKEND": "cuda"}, + Deployment("primary-only-ssd", DeviceEndpoint("hip", 0), None), + "auto", + 500, + None, + None, + False, + ) + 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"), + Deployment("strix", DeviceEndpoint("hip", 0), None), + 8080, + 8192, + [], + ) + self.assertIn("hip:0", 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) + + 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() 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 ab34f2a19..a24fa88ef 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -13,10 +13,12 @@ #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" #include +#include #include #include @@ -48,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 @@ -86,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); @@ -94,6 +98,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 @@ -106,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) { @@ -129,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; @@ -150,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(), @@ -162,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; } @@ -200,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()) { @@ -213,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(), @@ -412,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()) { @@ -436,6 +473,20 @@ 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; + cfg.moe_storage = plan.moe_storage_policy(); + + 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/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/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/feature_gate.cpp b/server/src/common/feature_gate.cpp index 0db7480d2..2d9dce0fc 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -146,6 +146,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/gpu_runtime_compat.h b/server/src/common/gpu_runtime_compat.h index dba8eaa7c..a9f14b6be 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/model_capabilities.h b/server/src/common/model_capabilities.h index 2189fe312..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,13 +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}, +// 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 = @@ -106,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() { @@ -235,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_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_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/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..797ad26eb 100644 --- a/server/src/common/moe_hybrid_storage.h +++ b/server/src/common/moe_hybrid_storage.h @@ -16,10 +16,27 @@ 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; +}; + +// 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). @@ -33,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. @@ -211,7 +229,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 +307,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..cef2a036f 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -1,13 +1,23 @@ -// 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 +#include +#include +#include #include +#include #include +#include +#include +#include +#include +#include +#include +#include +#include #if !defined(_WIN32) #include @@ -15,258 +25,1909 @@ #endif namespace dflash::common { +namespace { -MoeHybridStreamEngine::~MoeHybridStreamEngine() { - destroy(); +// 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; + } + + 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); + 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); +} + +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; +} + +} // 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.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.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( + "DFLASH_MOE_NVME_PREFILL_THRESHOLD", config.prefill_threshold, 1, 4096); + 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, + 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"); } -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; +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; } -MoeHybridStreamEngine & MoeHybridStreamEngine::operator=(MoeHybridStreamEngine && o) noexcept { - if (this != &o) { +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(); } + + 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, + size_t gate_alloc_bytes, + size_t up_alloc_bytes, + size_t down_alloc_bytes, + std::string * err) { 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; + 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_); + } + + if (spec.fused_gate_up) { + 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_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; + } + + 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_); + 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; +}; + +// 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 { + 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; + ggml_tensor * transfer_tensor = nullptr; + ggml_backend_event_t ready = nullptr; + bool pending = false; + bool valid = false; + bool cache_managed = false; + bool pinned = false; + int compute_users = 0; + MoeExpertKey key{}; + uint64_t frequency = 0; + uint64_t last_touch = 0; + MoeNvmeLease host_lease; + MoeExpertIoLayout layout{}; + DeviceExpertLayout device_layout{}; + }; + + ggml_backend_t backend = nullptr; + // 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; + // 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; + ggml_context * device_slot_ctx = nullptr; + size_t device_stride = 0; + 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; + uint64_t device_cache_evictions = 0; + 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; +}; + +template +void release_device_cache(RuntimeT & runtime) { + if (runtime.backend) ggml_backend_synchronize(runtime.backend); + 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) ggml_backend_event_free(slot.ready); + slot.ready = nullptr; + slot.data = nullptr; + slot.transfer_tensor = nullptr; + slot.pending = false; + } + 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); + } + 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; +} + +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) { + 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; + } + + 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; + 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; + 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; + } + if (attempt_slots == 2) break; + attempt_slots = std::max(2, attempt_slots * 3 / 4); + } + if (!runtime.device_pool) { + if (err) *err = "failed to allocate SSD GPU expert cache"; + return false; + } + + try { + runtime.device_slots.resize(attempt_slots); + } catch (const std::bad_alloc &) { + 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"; + 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; + } + for (size_t i = 0; i < runtime.device_slots.size(); ++i) { + 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; +} + +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; +} + +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) { + 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"; 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; + 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; + 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"; + ggml_backend_free(runtime->transfer_backend); + runtime->transfer_backend = nullptr; return false; } - pinned_size_ = max_expert_bytes; + if (!runtime->io->init(runtime->config.nvme, max_expert_bytes, + 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)) { + // Install the partial runtime so destroy() releases every resource. + runtime_ = std::move(runtime); + destroy(); + return false; + } + runtime_ = std::move(runtime); + return true; +} - // 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; +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) { + if (!init(gpu_backend, max_expert_bytes, config, err)) return false; + 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_backend && + !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; + if (runtime_->backend) ggml_backend_synchronize(runtime_->backend); + 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) ggml_backend_event_free(slot.ready); + slot.ready = nullptr; + slot.data = nullptr; + slot.transfer_tensor = nullptr; + slot.pending = false; + } + runtime_->device_slots.clear(); + runtime_->device_index.clear(); + if (runtime_->device_pool_buffer) { + ggml_backend_buffer_free(runtime_->device_pool_buffer); + } + 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; + 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 " + "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 " + "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, + 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.demand_timeouts, + (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, + (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.fused_decode_launches, + (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; + } + 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"; + 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.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; + } + if (dst.pending) { + ggml_backend_event_synchronize(dst.ready); + 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.pinned = false; + dst.key = {}; + dst.layout = MoeExpertIoLayout{}; + dst.device_layout = Runtime::DeviceExpertLayout{}; - const auto * file_base = static_cast(mmap_data); - size_t staging_offset = 0; - - // Validate expert_id against region size - if (expert_id < 0) { - if (err) *err = "expert_id is negative"; + if (!dst.ready) { + 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; } - // 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"; + 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; + } + 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; } - 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"; + 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]; + 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; } - 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"; + + 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; } - std::memcpy(static_cast(pinned_buf_) + staging_offset, - file_base + file_off, bytes); - last_up_bytes_ = bytes; - staging_offset += bytes; + + ggml_backend_tensor_set_async( + 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 { + // 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]; + ggml_backend_tensor_set_async( + runtime_->transfer_backend, dst.transfer_tensor, + lease.data() + span.buffer_offset, + span.device_offset, span.bytes); + } + 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}; + } + } + ggml_backend_event_record(dst.ready, runtime_->transfer_backend); + 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; + } } + runtime_->device_index.erase(cached); } - // 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"; + ++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.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; + const uint64_t score = (slot.frequency << 16) | recency; + if (score < best_score) { + best_score = score; + victim = (int) i; + } + } + } + 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; +} + +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; + } + Runtime::DeviceSlot & slot = runtime_->device_slots[(size_t) device_slot]; + if (slot.pending) { + ggml_backend_event_synchronize(slot.ready); + slot.pending = false; + slot.host_lease.reset(); + } + if (slot.layout.component_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; } - std::memcpy(static_cast(pinned_buf_) + staging_offset, - file_base + file_off, bytes); - last_down_bytes_ = bytes; - staging_offset += bytes; + ++slot.compute_users; + ++slot.frequency; + slot.last_touch = ++runtime_->device_clock; } + runtime_->active_slot = device_slot; + return true; +} + +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; +} + +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; +} + +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); + + 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; + } - if (staging_offset > scratch_size_) { - if (err) *err = "expert exceeds scratch buffer size"; + 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; +} + +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; } + 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); +} + +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]; + const MoeExpertComponentKind kind = slot.layout.fused_gate_up + ? MoeExpertComponentKind::FusedGateUp : MoeExpertComponentKind::Gate; + const Runtime::DeviceComponentLayout * component = + slot.device_layout.component(kind); + return component + ? static_cast(slot.data) + component->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; + const Runtime::DeviceComponentLayout * component = + slot.device_layout.component(MoeExpertComponentKind::Up); + return component + ? 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 Runtime::DeviceComponentLayout * component = + slot.device_layout.component(MoeExpertComponentKind::Down); + return component + ? static_cast(slot.data) + component->offset + : nullptr; +} + +size_t MoeHybridStreamEngine::scratch_gate_bytes() const { + if (!runtime_ || runtime_->active_slot < 0) return 0; + 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 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 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 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 { + return runtime_ && runtime_->io ? runtime_->io->total_host_bytes() : 0; +} - // 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; - } +size_t MoeHybridStreamEngine::scratch_bytes() const { + return runtime_ ? runtime_->device_pool_bytes : 0; +} - // 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_; - } - scratch_down_ = scratch_bytes + off; +const char * MoeHybridStreamEngine::io_backend_name() const { + return runtime_ && runtime_->io ? runtime_->io->effective_backend_name() : "uninitialized"; +} - return true; +MoeNvmeStats MoeHybridStreamEngine::io_stats() const { + return runtime_ && runtime_->io ? runtime_->io->stats() : MoeNvmeStats{}; } -// ── Streaming prefill evaluation ──────────────────────────────────────────── +MoeStreamComputeStats MoeHybridStreamEngine::compute_stats() const { + return runtime_ ? runtime_->compute_stats : MoeStreamComputeStats{}; +} -bool eval_moe_cold_experts_streaming( +static bool eval_moe_cold_experts_streaming_reference( MoeHybridStreamEngine & engine, ggml_backend_t gpu_backend, const void * mmap_data, @@ -280,155 +1941,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 +2143,960 @@ 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; +} + +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); + 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, + (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); + + // 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; + } + + // 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, execution_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; + 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); + 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(), + gate_component->alloc_bytes, + up_component ? up_component->alloc_bytes : 0, + down_component->alloc_bytes, + 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; + }; + + 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 < execution_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); + }; + + hits.clear(); + 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) { + 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; + } + 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; + 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 < execution_experts.size()) { + int next_slot = -1; + if (!engine.stage_expert_cached_async( + batch.layer, execution_experts[expert_index + 1], + &next_slot, err)) { + ggml_backend_synchronize(runtime.backend); + release_current(); + return false; + } + staged_slot = next_slot; + } + + if (!graph->finish(result, err)) { + 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 = 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; + for (int j = 0; j < spec.output_dim; ++j) { + dst[j] += weight * src[(size_t) j]; + } + } + 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; +} + +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 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, + 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; + 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; + 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; } -} // namespace dflash::common +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, + 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 d5e929ae9..733728e97 100644 --- a/server/src/common/moe_hybrid_stream.h +++ b/server/src/common/moe_hybrid_stream.h @@ -1,34 +1,157 @@ -// 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 + // 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; + // 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; + // 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. + size_t device_cache_bytes = 0; + MoeNvmeConfig nvme{}; + + 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; + 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 { + 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 +// 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; }; -// Streaming engine: manages pinned staging buffer, GPU scratch, and DMA pipeline. class MoeHybridStreamEngine { public: - MoeHybridStreamEngine() = default; + MoeHybridStreamEngine(); ~MoeHybridStreamEngine(); MoeHybridStreamEngine(const MoeHybridStreamEngine &) = delete; @@ -36,63 +159,174 @@ 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 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); + 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; + 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); + + // 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; + + size_t pinned_bytes() const; + 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_; +}; + +// 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(); - // Total pinned buffer size. - size_t pinned_bytes() const { return pinned_size_; } - // Total GPU scratch size. - size_t scratch_bytes() const { return scratch_size_; } + bool eval(const MoeStreamExpertSpec & spec, + const MoeStreamRouteBatch & batch, + const MoeStreamDualOwnerPolicy & policy, + std::vector & out, + MoeStreamDualOwnerStats * stats = nullptr, + std::string * err = nullptr); 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 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); + +// 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); + +// 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. bool eval_moe_cold_experts_streaming( MoeHybridStreamEngine & engine, ggml_backend_t gpu_backend, @@ -107,6 +341,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..8b46cfd14 --- /dev/null +++ b/server/src/common/moe_nvme_scheduler.cpp @@ -0,0 +1,1782 @@ +// 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 +#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_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 < 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) { + 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; + } + bytes_read = done; + 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); + 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); + 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 = MoeExpertIoLayout{}; + 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; + span.io_required_bytes = raw_io_bytes; + return true; + }; + + 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") || + !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; + } + } + + 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 demand_timeouts{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 (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; + } + 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 required_bytes = direct_active + ? span.io_required_bytes : span.bytes; + const size_t buffer_offset = direct_active ? span.io_buffer_offset : span.buffer_offset; + 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 + } + } + 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; + uint32_t required = 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 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 > (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 supported 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, + (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; + } + } + + // 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 < 0) { + item.ok = false; + 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); + } + } + 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_ = MoeExpertIoLayout{}; +} + +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)); + 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"; + 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 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"; + 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(); + 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) { + 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) { + if (!wait_for_state()) return false; + 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) { + if (!wait_for_state()) return false; + 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.demand_timeouts = p.demand_timeouts.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.demand_timeouts.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..d931a4180 --- /dev/null +++ b/server/src/common/moe_nvme_scheduler.h @@ -0,0 +1,265 @@ +// 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; + + // 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 + // 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 + // 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 { + 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. +// 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 demand_timeouts = 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/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/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/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 1451a23a2..23a671e76 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -524,6 +524,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; @@ -606,8 +642,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 = + cfg_.moe_storage == MoeStoragePolicy::Ssd; 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, @@ -893,8 +931,12 @@ bool DeepSeek4Backend::init_moe_tensor_parallel() { const Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); if (tp.in_process) { - 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; @@ -904,9 +946,10 @@ bool DeepSeek4Backend::init_moe_tensor_parallel() { cfg_.device.backend == PlacementBackend::Auto ? compiled_placement_backend() : cfg_.device.backend; std::fprintf(stderr, - "[deepseek4-moe-tp] enabled mode=in-process local=%s:%d " + "[deepseek4-moe-tp] enabled mode=%s local=%s:%d " "secondary=%s:%d primary_experts=%d " "secondary_experts=%d\n", + streamed_ready ? "in-process+ssd" : "in-process", placement_backend_name(local_kind), cfg_.device.gpu, placement_backend_name(tp.secondary_backend), tp.secondary_gpu, @@ -957,6 +1000,18 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & const Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); int hot_per_layer = tp.all_on_secondary ? 0 : budget.max_hot_per_layer; + if (!tp.all_on_secondary && + cfg_.moe_storage == MoeStoragePolicy::Ssd && + hot_per_layer >= w.n_expert) { + // `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 + // 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 (tp.all_on_secondary) { std::fprintf(stderr, "[deepseek4-moe-tp] all routed experts assigned to the " @@ -1067,8 +1122,17 @@ 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 Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); const bool inprocess_tp = tp.requested && tp.in_process; + const bool external_tp = tp.requested && !inprocess_tp; + // 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. `resident` explicitly selects the older materialized CPU + // tail. + bool stream_cold = !external_tp && + cfg_.moe_storage != MoeStoragePolicy::Resident; if (inprocess_tp) { const int expert_gpu = tp.secondary_gpu; const PlacementBackend expert_kind = tp.secondary_backend; @@ -1110,9 +1174,37 @@ bool DeepSeek4Backend::init_hybrid_model() { err.c_str()); return false; } - hybrid_cfg.materialize_cold_experts = true; hybrid_cfg.cold_expert_backend = MoeHybridColdBackend::Gpu; + + // On a separate expert GPU, auto mode retains the established fully + // resident path whenever the cold stack fits after a conservative + // reserve. Explicit ssd/resident modes remain authoritative. + if (cfg_.moe_storage == MoeStoragePolicy::Auto) { + Ds4ExpertMemoryInfo info; + std::string memory_error; + size_t expert_free = 0; + size_t expert_total = 0; + if (ggml_backend_dev_t device = + ggml_backend_get_device(expert_backend_)) { + ggml_backend_dev_memory(device, &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] 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), + stream_cold ? "ssd-stream" : "resident"); + } + } } + 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_)) { @@ -1124,34 +1216,74 @@ 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) { 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 " + "device=%d device_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); + inprocess_tp ? tp.secondary_gpu : cfg_.device.gpu, + stream_engine_.device_cache_bytes() / 1024.0 / 1024.0, + stream_engine_.device_slot_count()); } 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; @@ -1398,7 +1530,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, @@ -1530,7 +1663,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_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_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 9228d6387..b0127b77c 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -2981,7 +2981,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_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/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/kimi_k3/kimi_k3_backend.cpp b/server/src/kimi_k3/kimi_k3_backend.cpp new file mode 100644 index 000000000..c34bb7715 --- /dev/null +++ b/server/src/kimi_k3/kimi_k3_backend.cpp @@ -0,0 +1,946 @@ +#if defined(_WIN32) && !defined(NOMINMAX) +#define NOMINMAX +#endif + +#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" +#include "common/sampler.h" +#include "dflash27b.h" + +#include "ggml-cuda.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#include +#include +#include +#include +#else +#include +#include +#include +#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) {} + +KimiK3Backend::~KimiK3Backend() { + shutdown(); +} + +void KimiK3Backend::release_expert_backend() { + if (expert_backend_) { + ggml_backend_free(expert_backend_); + expert_backend_ = nullptr; + } + expert_gpu_ = -1; + expert_backend_kind_ = PlacementBackend::Auto; +} + +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) { + std::fprintf(stderr, + "[kimi-k3] routed expert streaming metadata is incomplete\n"); + 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; + 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 %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(); + stream_engine_.destroy(); + secondary_stream_engine_.destroy(); + stream_owner_policy_ = MoeStreamDualOwnerPolicy{}; + stream_placement_ = MoeHybridPlacement{}; + release_expert_backend(); + return false; + }; + + 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); + } + 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 = + std::min(stream_config.device_cache_bytes, routed_pool_bytes); + return stream_config; + } + size_t free_bytes = 0; + size_t total_bytes = 0; + if (ggml_backend_dev_t device = + 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 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 " + "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(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; + }; + + stream_owner_policy_ = MoeStreamDualOwnerPolicy::from_env(); + stream_placement_ = MoeHybridPlacement{}; + const char * placement_path = std::getenv("DFLASH_MOE_PLACEMENT"); + 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 streamed-expert placement %s: %s\n", + placement_path, + error.empty() ? "model shape mismatch" : error.c_str()); + return fail_streaming(); + } + if (expert_backend_) { + stream_owner_policy_.primary_placement = &stream_placement_; + } + } + + ScopedFileDescriptors descriptors; + std::vector model_sources; + model_sources.reserve(weights_.shard_paths.size()); + for (const std::string & shard : weights_.shard_paths) { + MoeNvmeSource source; + if (!open_nvme_source(shard, descriptors, source, error)) { + std::fprintf(stderr, + "[kimi-k3] %s\n", error.c_str()); + return fail_streaming(); + } + 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(); + } + 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] 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( + backend_, cfg_.device.primary_gpu(), primary_kind, "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_backend_, expert_gpu_, expert_backend_kind_, "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(); + } + } + + const bool primary_bound = stream_engine_.bind_sources( + *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( + *active_sources, *active_regions, &secondary_error); + } + 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(); + } + + 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, + "[kimi-k3] dual-owner executor initialization failed: %s\n", + 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 " + "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(), + 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), + 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, + 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; +} + +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; + } + const bool stream_routed_experts = + cfg_.moe_storage != MoeStoragePolicy::Resident; + if (!load_kimi_k3_gguf( + cfg_.model_path, backend_, weights_, + stream_routed_experts)) { + 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; + } + 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", + cfg_.device.primary_gpu(), max_ctx, + !weights_.routed_experts_streamed ? "resident" : + (expert_backend_ ? "nvme-dual-owner" : "nvme-single-owner")); + 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_) { + maybe_save_routing_stats(); + dual_stream_executor_.destroy(); + stream_engine_.destroy(); + secondary_stream_engine_.destroy(); + release_expert_backend(); + 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_, + cfg_.moe_storage != MoeStoragePolicy::Resident) || + (weights_.routed_experts_streamed && !init_streaming())) { + 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, &stream_engine_, + dual_stream_executor_.is_ready() + ? &dual_stream_executor_ : nullptr, + &stream_owner_policy_, routing_stats_.get())) { + 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, &stream_engine_, + dual_stream_executor_.is_ready() + ? &dual_stream_executor_ : nullptr, + &stream_owner_policy_, routing_stats_.get())) { + 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(); + maybe_save_routing_stats(); + 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() { + maybe_save_routing_stats(); + dual_stream_executor_.destroy(); + stream_engine_.destroy(); + secondary_stream_engine_.destroy(); + release_expert_backend(); + free_kimi_k3_cache(cache_); + free_kimi_k3_weights(weights_); + if (backend_) { + 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 new file mode 100644 index 000000000..f5472e992 --- /dev/null +++ b/server/src/kimi_k3/kimi_k3_backend.h @@ -0,0 +1,86 @@ +#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 { + +struct KimiK3BackendConfig { + const char * model_path = nullptr; + DevicePlacement device; + int stream_fd = -1; + // -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. + MoeStoragePolicy moe_storage = MoeStoragePolicy::Auto; +}; + +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: + bool init_streaming(); + void release_expert_backend(); + void maybe_save_routing_stats(); + + int32_t choose_token(const std::vector & logits, + const SamplerCfg & sampler, + const std::vector & history); + + 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_; + MoeHybridStreamEngine stream_engine_; + MoeHybridStreamEngine secondary_stream_engine_; + 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{}()}; +}; + +} // 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..38788a4b4 --- /dev/null +++ b/server/src/kimi_k3/kimi_k3_graph.cpp @@ -0,0 +1,942 @@ +#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" + +#include "ggml-alloc.h" + +#include +#include +#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); +} + +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, + 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); + } + 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); + 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); +} + +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, + MoeStreamDualOwnerExecutor * dual_stream_executor, + const MoeStreamDualOwnerPolicy * stream_owner_policy, + MoeHybridRoutingStats * routing_stats) { + 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(); + 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; + 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); + if (!route_ok) { + set_last_error( + "Kimi-K3 routed layer " + + std::to_string(il) + + ": streamed expert evaluation failed: " + + 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) { + 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, + 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, + MoeHybridStreamEngine * stream_engine, + MoeStreamDualOwnerExecutor * dual_stream_executor, + 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) { + 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; + } + 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, dual_stream_executor, + stream_owner_policy, routing_stats); + } + + 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..42d9f6cf7 --- /dev/null +++ b/server/src/kimi_k3/kimi_k3_internal.h @@ -0,0 +1,184 @@ +// 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 "common/moe_hybrid_storage.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include + +namespace dflash::common { + +struct MoeHybridRoutingStats; + +class MoeHybridStreamEngine; +class MoeStreamDualOwnerExecutor; +struct MoeStreamDualOwnerPolicy; + +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 { + // 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; + 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, + bool stream_routed_experts = false); +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, + MoeHybridStreamEngine * stream_engine = nullptr, + MoeStreamDualOwnerExecutor * dual_stream_executor = nullptr, + const MoeStreamDualOwnerPolicy * stream_owner_policy = nullptr, + MoeHybridRoutingStats * routing_stats = nullptr); + +} // 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..2919e9293 --- /dev/null +++ b/server/src/kimi_k3/kimi_k3_loader.cpp @@ -0,0 +1,622 @@ +#include "kimi_k3_internal.h" + +#include "common/gguf_bounds.h" +#include "common/gguf_mmap.h" +#include "internal.h" + +#include +#include +#include +#include +#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_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; +} + +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; +} + +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, + bool stream_routed_experts) { + free_kimi_k3_weights(out); + + 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 (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); + for (gguf_context * gguf : shard_ggufs) gguf_free(gguf); + shard_ggufs.clear(); + free_kimi_k3_weights(out); + return false; + }; + + 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) { + 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 = 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")); + 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)); + + 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"); + 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.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); + } + + 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; + 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; + } + } + out.buf = out.buffers.empty() ? nullptr : out.buffers.front(); + + for (gguf_context * gguf : shard_ggufs) gguf_free(gguf); + shard_ggufs.clear(); + std::fprintf(stderr, + "[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; })), + 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) { + 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{}; +} + +} // namespace dflash::common 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/src/server/server_main.cpp b/server/src/server/server_main.cpp index 9cb7feafc..ff54c75c8 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -101,6 +101,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" @@ -325,6 +327,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) { @@ -975,6 +995,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())); std::fprintf(stderr, "[server] │ target_split = %s\n", target_split_mode_name(bargs.device.split_mode)); if (bargs.device.is_multi_device()) { diff --git a/server/test/bench_kimi_k3_hetero.cpp b/server/test/bench_kimi_k3_hetero.cpp new file mode 100644 index 000000000..a1269cc65 --- /dev/null +++ b/server/test/bench_kimi_k3_hetero.cpp @@ -0,0 +1,499 @@ +// Read-only Kimi-K3 routed-core qualification for heterogeneous Lucebox. +// +// 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 +#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); +} + +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, n_expert - 1); + const int selected = choose(rng); + std::swap(population[(size_t) i], population[(size_t) selected]); + } + population.resize((size_t) top_k); + 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=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 = 0; + uint64_t top_k_arg = 0; + 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 || 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; + + 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; + } + + 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) { + 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; + } + + MoeStreamExpertSpec expert_spec; + 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 = situ_beta; + expert_spec.situ_linear_beta = situ_linear_beta; + + 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); + } + 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; + 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, n_expert, (int) top_k_arg, repeat_arg != 0); + if (compute_arg) { + MoeStreamRouteBatch batch; + batch.layer = layer; + batch.n_expert = n_expert; + 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); + + 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; + } + // 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( + 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; + } + 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 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 + ? 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) 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", + 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 + " graph_builds=%" PRIu64 " graph_hits=%" PRIu64 + " 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.fused_decode_launches, + compute_stats.fused_decode_experts); + + 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/smoke_kimi_k3_forward.cpp b/server/test/smoke_kimi_k3_forward.cpp new file mode 100644 index 000000000..c635a84c8 --- /dev/null +++ b/server/test/smoke_kimi_k3_forward.cpp @@ -0,0 +1,75 @@ +#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] " + "[stream_experts=1] [expert_gpu=-1]\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; + 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; + + 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 7c34188dd..9f83ee650 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"; @@ -332,9 +369,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()); } @@ -343,6 +380,39 @@ 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()); +} + +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 ───────────────────────────────────────────────── @@ -383,9 +453,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")); @@ -447,7 +518,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("")); @@ -459,6 +530,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. @@ -470,10 +542,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); @@ -487,6 +566,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); 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_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_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_nvme_scheduler.cpp b/server/test/test_moe_nvme_scheduler.cpp new file mode 100644 index 000000000..ccb90356f --- /dev/null +++ b/server/test/test_moe_nvme_scheduler.cpp @@ -0,0 +1,545 @@ +#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); + 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; + 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.component_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, 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; + 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, 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); + 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, 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); + 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; +#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; + 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)); +#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); + lease.reset(); + NVME_REQUIRE(scheduler.stats().read_ops == 2); + NVME_REQUIRE(scheduler.stats().errors == 0); + scheduler.destroy(); + ::close(fd); +} +#endif diff --git a/server/test/test_moe_stream_compute.cpp b/server/test/test_moe_stream_compute.cpp new file mode 100644 index 000000000..c4980a099 --- /dev/null +++ b/server/test/test_moe_stream_compute.cpp @@ -0,0 +1,686 @@ +#include "CppUnitTestFramework.hpp" +#include "common/dynamic_backend.h" +#include "common/moe_hybrid_stream.h" + +#include "ggml-quants.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 = 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. +constexpr int kInput = 256; +constexpr int kFf = 64; +constexpr int kOutput = 128; +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 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, + const std::vector & down, + const std::vector & input, + const int32_t * ids, + 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) n_tokens * kOutput, 0.0f); + std::vector activated(kFf); + 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; + 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 * 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); + + // 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); + 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); + 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) { + 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; + 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(); +} + +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(); +} + +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) { + int device = 0; + if (const char * value = std::getenv("DFLASH_TEST_GPU")) { + device = std::max(0, std::atoi(value)); + } + 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); + } + std::string init_error; + ggml_backend_t backend = init_placement_backend( + backend_kind, device, &init_error); + if (!backend) { + 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); + 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); +} 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..2862f3623 --- /dev/null +++ b/server/test/test_moe_stream_owner_partition.cpp @@ -0,0 +1,189 @@ +#include "CppUnitTestFramework.hpp" +#include "../src/common/moe_hybrid_stream.h" +#include "../src/common/moe_stream_cache_policy.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)); +} + +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)); + } +} 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..3b4583bd5 100644 --- a/variables.md +++ b/variables.md @@ -217,6 +217,9 @@ Untagged variables are operational tuning knobs. | Variable | Purpose | |---|---| +| `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. |