From bfacd04c0cc3f6af21603d8edca433be144f7839 Mon Sep 17 00:00:00 2001 From: weicj Date: Fri, 3 Jul 2026 01:22:36 +0800 Subject: [PATCH 1/2] feat(server): add MoE expert-split foundation --- server/CMakeLists.txt | 59 +- server/src/common/backend_ipc.cpp | 54 +- server/src/common/backend_ipc.h | 10 + server/src/common/cold_ffn_compute.h | 62 +- .../common/expert_split_compute_runtime.cpp | 122 + .../src/common/expert_split_compute_runtime.h | 44 + .../common/expert_split_materialization.cpp | 166 ++ .../src/common/expert_split_materialization.h | 40 + server/src/common/expert_split_plan.cpp | 307 +++ server/src/common/expert_split_plan.h | 76 + server/src/common/expert_split_runtime.cpp | 279 +++ server/src/common/expert_split_runtime.h | 72 + server/src/common/expert_split_state.cpp | 124 + server/src/common/expert_split_state.h | 65 + .../src/common/expert_split_target_config.cpp | 347 +++ .../src/common/expert_split_target_config.h | 52 + server/src/common/gguf_tensor_data.cpp | 480 ++++ server/src/common/gguf_tensor_data.h | 76 + server/src/common/moe_expert_compute.cpp | 1284 ++++++++++ server/src/common/moe_expert_compute.h | 261 +++ server/src/common/moe_expert_compute_cpu.cpp | 195 ++ server/src/common/moe_expert_compute_ipc.cpp | 2075 +++++++++++++++++ server/src/common/moe_hybrid_ffn_eval.cpp | 1253 ++++++++-- server/src/common/moe_hybrid_ffn_eval.h | 8 +- server/src/common/moe_hybrid_placement.cpp | 136 +- server/src/common/moe_hybrid_storage.cpp | 189 +- server/src/common/moe_hybrid_storage.h | 26 +- server/src/internal.h | 44 + server/src/ipc/backend_ipc_main.cpp | 16 +- server/src/qwen35/gguf_target_loader.cpp | 208 +- server/src/qwen35/graph_builders.cpp | 6 +- server/src/qwen35/qwen35_backend.cpp | 4 +- server/src/qwen35/qwen35_target_graph.cpp | 45 +- server/src/qwen35moe/qwen35moe_backend.cpp | 378 ++- server/src/qwen35moe/qwen35moe_backend.h | 29 +- .../qwen35moe/qwen35moe_pipelined_decode.cpp | 336 ++- .../qwen35moe/qwen35moe_pipelined_decode.h | 51 +- .../test_expert_split_materialization.cpp | 104 + server/test/test_expert_split_plan.cpp | 81 + server/test/test_expert_split_runtime.cpp | 257 ++ server/test/test_expert_split_state.cpp | 86 + .../test/test_expert_split_target_config.cpp | 113 + .../test_moe_expert_compute_multi_target.cpp | 354 +++ server/test/test_server_unit.cpp | 41 + 44 files changed, 9275 insertions(+), 740 deletions(-) create mode 100644 server/src/common/expert_split_compute_runtime.cpp create mode 100644 server/src/common/expert_split_compute_runtime.h create mode 100644 server/src/common/expert_split_materialization.cpp create mode 100644 server/src/common/expert_split_materialization.h create mode 100644 server/src/common/expert_split_plan.cpp create mode 100644 server/src/common/expert_split_plan.h create mode 100644 server/src/common/expert_split_runtime.cpp create mode 100644 server/src/common/expert_split_runtime.h create mode 100644 server/src/common/expert_split_state.cpp create mode 100644 server/src/common/expert_split_state.h create mode 100644 server/src/common/expert_split_target_config.cpp create mode 100644 server/src/common/expert_split_target_config.h create mode 100644 server/src/common/gguf_tensor_data.cpp create mode 100644 server/src/common/gguf_tensor_data.h create mode 100644 server/src/common/moe_expert_compute.cpp create mode 100644 server/src/common/moe_expert_compute.h create mode 100644 server/src/common/moe_expert_compute_cpu.cpp create mode 100644 server/src/common/moe_expert_compute_ipc.cpp create mode 100644 server/test/test_expert_split_materialization.cpp create mode 100644 server/test/test_expert_split_plan.cpp create mode 100644 server/test/test_expert_split_runtime.cpp create mode 100644 server/test/test_expert_split_state.cpp create mode 100644 server/test/test_expert_split_target_config.cpp create mode 100644 server/test/test_moe_expert_compute_multi_target.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 1b894c0af..4caffc41f 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -269,6 +269,13 @@ add_library(dflash_common STATIC src/common/dflash_spec_decode.cpp src/common/layer_split_backend.cpp src/common/layer_split_runtime.cpp + src/common/expert_split_plan.cpp + src/common/expert_split_state.cpp + src/common/expert_split_runtime.cpp + src/common/expert_split_compute_runtime.cpp + src/common/expert_split_materialization.cpp + src/common/expert_split_target_config.cpp + src/common/gguf_tensor_data.cpp src/qwen35/graph_builders.cpp src/qwen35moe/qwen35moe_ffn.cpp src/qwen35moe/qwen35moe_backend.cpp @@ -281,7 +288,9 @@ add_library(dflash_common STATIC src/common/spark_corpus.cpp src/common/moe_hybrid_ffn_eval.cpp src/common/moe_hybrid_stream.cpp - src/common/cold_ffn_cpu.cpp + src/common/moe_expert_compute.cpp + src/common/moe_expert_compute_cpu.cpp + src/common/moe_expert_compute_ipc.cpp src/common/moe_hybrid_swap_manager.cpp src/common/moe_routing_collector.cpp src/qwen35/layer_split_forward.cpp @@ -728,6 +737,54 @@ if(DFLASH27B_TESTS) endif() target_link_libraries(test_qwen35moe_swap_manager PRIVATE dflash_common) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_expert_split_plan.cpp") + add_executable(test_expert_split_plan test/test_expert_split_plan.cpp) + target_include_directories(test_expert_split_plan PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_include_directories(test_expert_split_plan PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + endif() + target_link_libraries(test_expert_split_plan PRIVATE dflash_common) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_expert_split_runtime.cpp") + add_executable(test_expert_split_runtime test/test_expert_split_runtime.cpp) + target_include_directories(test_expert_split_runtime PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_include_directories(test_expert_split_runtime PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + endif() + target_link_libraries(test_expert_split_runtime PRIVATE dflash_common) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_expert_split_state.cpp") + add_executable(test_expert_split_state test/test_expert_split_state.cpp) + target_include_directories(test_expert_split_state PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_include_directories(test_expert_split_state PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + endif() + target_link_libraries(test_expert_split_state PRIVATE dflash_common) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_expert_split_target_config.cpp") + add_executable(test_expert_split_target_config test/test_expert_split_target_config.cpp) + target_include_directories(test_expert_split_target_config PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_include_directories(test_expert_split_target_config PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + endif() + target_link_libraries(test_expert_split_target_config PRIVATE dflash_common) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_expert_split_materialization.cpp") + add_executable(test_expert_split_materialization test/test_expert_split_materialization.cpp) + target_include_directories(test_expert_split_materialization PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_include_directories(test_expert_split_materialization PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + endif() + target_link_libraries(test_expert_split_materialization PRIVATE dflash_common) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_moe_expert_compute_multi_target.cpp") + add_executable(test_moe_expert_compute_multi_target test/test_moe_expert_compute_multi_target.cpp) + target_include_directories(test_moe_expert_compute_multi_target PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_include_directories(test_moe_expert_compute_multi_target PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + endif() + target_link_libraries(test_moe_expert_compute_multi_target PRIVATE dflash_common) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/smoke_load_draft.cpp") add_executable(smoke_load_draft test/smoke_load_draft.cpp) target_include_directories(smoke_load_draft PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) diff --git a/server/src/common/backend_ipc.cpp b/server/src/common/backend_ipc.cpp index 0e4d38c4e..3633d20fd 100644 --- a/server/src/common/backend_ipc.cpp +++ b/server/src/common/backend_ipc.cpp @@ -32,6 +32,7 @@ const char * backend_ipc_mode_name(BackendIpcMode mode) { case BackendIpcMode::Qwen35TargetShard: return "qwen35-target-shard"; case BackendIpcMode::Gemma4TargetShard: return "gemma4-target-shard"; case BackendIpcMode::LagunaTargetShard: return "laguna-target-shard"; + case BackendIpcMode::MoeExpertCompute: return "moe-expert-compute"; } return "unknown"; } @@ -57,6 +58,10 @@ bool parse_backend_ipc_mode(const std::string & value, BackendIpcMode & out) { out = BackendIpcMode::LagunaTargetShard; return true; } + if (value == "moe-expert-compute") { + out = BackendIpcMode::MoeExpertCompute; + return true; + } return false; } @@ -257,13 +262,33 @@ std::string BackendIpcProcess::next_path(const char * prefix) { } bool BackendIpcProcess::write_shared_payload(const void * data, size_t bytes, uint64_t & seq) { - if (!shared_payload_map_ || bytes > shared_payload_capacity_) return false; - if (bytes > 0 && !data) return false; + BackendIpcPayloadSegment segment{data, bytes}; + return write_shared_payload_segments(&segment, 1, seq); +} + +bool BackendIpcProcess::write_shared_payload_segments( + const BackendIpcPayloadSegment * segments, + size_t n_segments, + uint64_t & seq) { + if (!shared_payload_map_ || (!segments && n_segments > 0)) return false; + size_t bytes = 0; + for (size_t i = 0; i < n_segments; ++i) { + if (segments[i].bytes > 0 && !segments[i].data) return false; + if (!backend_ipc_checked_add_size(bytes, segments[i].bytes, bytes)) { + return false; + } + } + if (bytes > shared_payload_capacity_) return false; auto * header = static_cast(shared_payload_map_); - void * payload = static_cast( - static_cast(shared_payload_map_) + backend_ipc_shared_payload_header_bytes()); - if (bytes > 0) { - std::memcpy(payload, data, bytes); + auto * payload = static_cast( + static_cast(shared_payload_map_) + + backend_ipc_shared_payload_header_bytes()); + size_t off = 0; + for (size_t i = 0; i < n_segments; ++i) { + if (segments[i].bytes > 0) { + std::memcpy(payload + off, segments[i].data, segments[i].bytes); + off += segments[i].bytes; + } } seq = ++shared_payload_seq_; header->bytes = (uint64_t)bytes; @@ -271,6 +296,23 @@ bool BackendIpcProcess::write_shared_payload(const void * data, size_t bytes, ui return true; } +bool BackendIpcProcess::read_shared_payload(void * data, size_t bytes, uint64_t seq) const { + if (!shared_payload_map_ || bytes > shared_payload_capacity_) return false; + if (bytes > 0 && !data) return false; + const auto * header = + static_cast(shared_payload_map_); + if (header->sequence != seq || header->bytes != (uint64_t)bytes) { + return false; + } + const void * payload = static_cast( + static_cast(shared_payload_map_) + + backend_ipc_shared_payload_header_bytes()); + if (bytes > 0) { + std::memcpy(data, payload, bytes); + } + return true; +} + #if !defined(_WIN32) bool BackendIpcProcess::init_shared_payload(size_t bytes) { if (bytes == 0) return false; diff --git a/server/src/common/backend_ipc.h b/server/src/common/backend_ipc.h index b95fbc33a..62ea6f493 100644 --- a/server/src/common/backend_ipc.h +++ b/server/src/common/backend_ipc.h @@ -27,6 +27,7 @@ enum class BackendIpcMode { Qwen35TargetShard, Gemma4TargetShard, LagunaTargetShard, + MoeExpertCompute, }; const char * backend_ipc_mode_name(BackendIpcMode mode); @@ -82,6 +83,11 @@ struct BackendIpcLaunchConfig { size_t shared_payload_bytes = 0; }; +struct BackendIpcPayloadSegment { + const void * data = nullptr; + size_t bytes = 0; +}; + class BackendIpcProcess { public: BackendIpcProcess() = default; @@ -105,6 +111,10 @@ class BackendIpcProcess { std::string next_path(const char * prefix); bool write_shared_payload(const void * data, size_t bytes, uint64_t & seq); + bool write_shared_payload_segments(const BackendIpcPayloadSegment * segments, + size_t n_segments, + uint64_t & seq); + bool read_shared_payload(void * data, size_t bytes, uint64_t seq) const; private: #if !defined(_WIN32) diff --git a/server/src/common/cold_ffn_compute.h b/server/src/common/cold_ffn_compute.h index f1d512dec..2cbbf9f82 100644 --- a/server/src/common/cold_ffn_compute.h +++ b/server/src/common/cold_ffn_compute.h @@ -1,59 +1,19 @@ -// ColdFfnCompute: Direct compute interface for cold expert FFN. -// Bypasses ggml graph dispatch overhead. Shared-memory model (CPU/Halo). +// Compatibility shim for the old cold-FFN name. +// +// New code should include moe_expert_compute.h directly. The aliases keep +// existing call sites and downstream branches buildable while the compute +// abstraction moves from "cold expert fallback" to neutral MoE expert compute. #pragma once -#include "ggml.h" -#include -#include +#include "moe_expert_compute.h" namespace dflash::common { -// Per-layer cold weight metadata — raw pointers into shared memory. -struct ColdFfnLayer { - const void * gate_up_data = nullptr; // fused [n_cold, n_ff*2, n_embd] quantized - const void * gate_data = nullptr; // separate gate [n_cold, n_ff, n_embd] - const void * up_data = nullptr; // separate up [n_cold, n_ff, n_embd] - const void * down_data = nullptr; // [n_cold, n_embd, n_ff] quantized +using ColdFfnLayer = MoeExpertLayer; +using ColdFfnCompute = MoeExpertCompute; - size_t gate_up_stride = 0; // bytes between experts in gate_up tensor - size_t gate_stride = 0; // bytes between experts in gate tensor - size_t up_stride = 0; // bytes between experts in up tensor - size_t down_stride = 0; // bytes between experts in down tensor - - ggml_type gate_up_type = GGML_TYPE_Q4_K; // type for fused gate_up - ggml_type gate_type = GGML_TYPE_Q4_K; // type for separate gate - ggml_type up_type = GGML_TYPE_Q4_K; // type for separate up - ggml_type down_type = GGML_TYPE_Q4_K; // type for down projection - bool fused_gate_up = false; // true if gate+up are fused - - // Scale factors (applied after matmul). 1.0 = no scaling. - float gate_up_scale = 1.0f; - float gate_scale = 1.0f; - float up_scale = 1.0f; - float down_scale = 1.0f; -}; - -// Abstract compute interface. Implementations: CPU (now), Halo (future). -struct ColdFfnCompute { - virtual ~ColdFfnCompute() = default; - - // Compute cold expert FFN contributions and accumulate into output. - // input: [n_embd] F32 — post-norm hidden state - // ids: [n_cold] I32 — local cold expert indices - // weights: [n_cold] F32 — routing weights for each cold expert - // output: [n_embd] F32 — accumulated weighted expert outputs (zeroed by callee) - virtual void compute( - const ColdFfnLayer & layer, - const float * input, - const int32_t * ids, - const float * weights, - int n_cold, - int n_embd, - int n_ff, - float * output) = 0; -}; - -// Create CPU-based fused cold FFN compute. -std::unique_ptr make_cpu_cold_ffn_compute(int n_ff_max); +inline std::unique_ptr make_cpu_cold_ffn_compute(int n_ff_max) { + return make_cpu_moe_expert_compute(n_ff_max); +} } // namespace dflash::common diff --git a/server/src/common/expert_split_compute_runtime.cpp b/server/src/common/expert_split_compute_runtime.cpp new file mode 100644 index 000000000..6e6a0f9ae --- /dev/null +++ b/server/src/common/expert_split_compute_runtime.cpp @@ -0,0 +1,122 @@ +#include "expert_split_compute_runtime.h" + +#include + +namespace dflash::common { + +namespace { + +uint64_t hash_u64(uint64_t h, uint64_t v) { + h ^= v + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2); + return h; +} + +uint64_t hash_string(uint64_t h, const std::string & s) { + for (unsigned char c : s) { + h = hash_u64(h, (uint64_t)c); + } + return h; +} + +} // namespace + +bool ExpertSplitComputeRuntime::matches(int n_layer_, int n_expert_, int n_expert_used_) const { + return n_layer == n_layer_ && + n_expert == n_expert_ && + n_expert_used == n_expert_used_ && + (int) target_index_by_global.size() == n_layer * n_expert && + (int) local_index_by_global.size() == n_layer * n_expert; +} + +int ExpertSplitComputeRuntime::index(int layer, int expert) const { + if (layer < 0 || layer >= n_layer || expert < 0 || expert >= n_expert) { + return -1; + } + return layer * n_expert + expert; +} + +int ExpertSplitComputeRuntime::target_index(int layer, int expert) const { + const int idx = index(layer, expert); + return idx >= 0 ? target_index_by_global[(size_t)idx] : -1; +} + +int ExpertSplitComputeRuntime::local_index(int layer, int expert) const { + const int idx = index(layer, expert); + return idx >= 0 ? local_index_by_global[(size_t)idx] : -1; +} + +bool build_expert_split_compute_runtime(const ExpertSplitRuntime & runtime, + int n_expert_used, + ExpertSplitComputeRuntime & out, + std::string * err) { + if (!runtime.matches(runtime.n_layer, runtime.n_expert) || + runtime.n_layer <= 0 || runtime.n_expert <= 0) { + if (err) *err = "expert split runtime not initialized"; + return false; + } + if (n_expert_used <= 0 || n_expert_used > runtime.n_expert) { + if (err) *err = "invalid n_expert_used for expert split compute runtime"; + return false; + } + + std::vector placements; + if (!build_all_expert_split_target_placements(runtime, n_expert_used, placements, err)) { + return false; + } + if (placements.size() != runtime.targets.size()) { + if (err) *err = "expert split target placement count mismatch"; + return false; + } + + ExpertSplitComputeRuntime compute_runtime; + compute_runtime.n_layer = runtime.n_layer; + compute_runtime.n_expert = runtime.n_expert; + compute_runtime.n_expert_used = n_expert_used; + compute_runtime.target_index_by_global = runtime.target_index_by_global; + compute_runtime.local_index_by_global = runtime.local_index_by_global; + compute_runtime.targets.reserve(runtime.targets.size()); + + for (size_t i = 0; i < runtime.targets.size(); ++i) { + ExpertSplitComputeTargetRuntime target_runtime; + target_runtime.target_index = (int) i; + target_runtime.target = runtime.targets[i].target; + target_runtime.placement = std::move(placements[i].placement); + compute_runtime.targets.push_back(std::move(target_runtime)); + } + + out = std::move(compute_runtime); + return true; +} + +uint64_t expert_split_compute_runtime_fingerprint( + const ExpertSplitComputeRuntime & runtime) { + uint64_t h = 1469598103934665603ULL; + h = hash_u64(h, (uint64_t)runtime.n_layer); + h = hash_u64(h, (uint64_t)runtime.n_expert); + h = hash_u64(h, (uint64_t)runtime.n_expert_used); + h = hash_u64(h, (uint64_t)runtime.targets.size()); + for (size_t ti = 0; ti < runtime.targets.size(); ++ti) { + const ExpertSplitComputeTargetRuntime & target = runtime.targets[ti]; + h = hash_u64(h, (uint64_t)ti); + h = hash_string(h, target.target.name); + h = hash_string(h, target.target.backend); + h = hash_u64(h, (uint64_t)(int64_t)target.target.device_id); + h = hash_u64(h, target.placement.n_layer); + h = hash_u64(h, target.placement.n_expert); + h = hash_u64(h, target.placement.n_expert_used); + h = hash_u64(h, target.placement.total_hot); + h = hash_u64(h, (uint64_t)target.placement.hot_counts.size()); + for (int hot_count : target.placement.hot_counts) { + h = hash_u64(h, (uint64_t)hot_count); + } + for (size_t il = 0; il < target.placement.hot_expert_ids.size(); ++il) { + h = hash_u64(h, (uint64_t)il); + for (int32_t expert : target.placement.hot_expert_ids[il]) { + h = hash_u64(h, (uint64_t)(uint32_t)expert); + } + } + } + return h; +} + +} // namespace dflash::common diff --git a/server/src/common/expert_split_compute_runtime.h b/server/src/common/expert_split_compute_runtime.h new file mode 100644 index 000000000..104f0fae6 --- /dev/null +++ b/server/src/common/expert_split_compute_runtime.h @@ -0,0 +1,44 @@ +// ExpertSplit compute runtime scaffolding. +// +// This layer is intentionally non-executing for now: it derives per-target +// expert placements and leaves target-local compute/storage construction to a +// later phase. The goal is to preserve the ordered multi-target contract +// without disturbing the existing hot/cold execution path. + +#pragma once + +#include "expert_split_runtime.h" + +#include + +namespace dflash::common { + +struct ExpertSplitComputeTargetRuntime { + int target_index = -1; + ExpertSplitTarget target; + MoeHybridPlacement placement; +}; + +struct ExpertSplitComputeRuntime { + int n_layer = 0; + int n_expert = 0; + int n_expert_used = 0; + std::vector targets; + std::vector target_index_by_global; + std::vector local_index_by_global; + + bool matches(int n_layer_, int n_expert_, int n_expert_used_) const; + int index(int layer, int expert) const; + int target_index(int layer, int expert) const; + int local_index(int layer, int expert) const; +}; + +bool build_expert_split_compute_runtime(const ExpertSplitRuntime & runtime, + int n_expert_used, + ExpertSplitComputeRuntime & out, + std::string * err = nullptr); + +uint64_t expert_split_compute_runtime_fingerprint( + const ExpertSplitComputeRuntime & runtime); + +} // namespace dflash::common diff --git a/server/src/common/expert_split_materialization.cpp b/server/src/common/expert_split_materialization.cpp new file mode 100644 index 000000000..6a027179a --- /dev/null +++ b/server/src/common/expert_split_materialization.cpp @@ -0,0 +1,166 @@ +#include "expert_split_materialization.h" + +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +bool validate_runtime(const ExpertSplitRuntime & runtime, + int n_expert_used, + std::string * err) { + if (!runtime.matches(runtime.n_layer, runtime.n_expert) || + runtime.n_layer <= 0 || runtime.n_expert <= 0) { + if (err) *err = "expert split runtime not initialized"; + return false; + } + if (runtime.targets.empty()) { + if (err) *err = "expert split runtime has no targets"; + return false; + } + if (n_expert_used <= 0 || n_expert_used > runtime.n_expert) { + if (err) *err = "invalid n_expert_used for expert split materialization"; + return false; + } + return true; +} + +} // namespace + +bool ExpertSplitMaterialization::matches(int n_layer_, + int n_expert_, + int n_expert_used_) const { + if (n_layer != n_layer_ || n_expert != n_expert_ || + n_expert_used != n_expert_used_) { + return false; + } + if (!primary_placement.matches(n_layer_, n_expert_, n_expert_used_)) { + return false; + } + if (ordered_cold_union && + (int) cold_expert_ids_by_layer.size() != n_layer_) { + return false; + } + if (!cold_expert_ids_by_layer.empty() && + (int) cold_expert_ids_by_layer.size() != n_layer_) { + return false; + } + return true; +} + +bool build_expert_split_materialization(const ExpertSplitRuntime & runtime, + int n_expert_used, + ExpertSplitMaterialization & out, + std::string * err) { + if (!validate_runtime(runtime, n_expert_used, err)) { + return false; + } + + std::vector placements; + if (!build_all_expert_split_target_placements(runtime, n_expert_used, + placements, err)) { + return false; + } + if (placements.size() != runtime.targets.size()) { + if (err) *err = "expert split target placement count mismatch"; + return false; + } + + ExpertSplitMaterialization materialization; + materialization.n_layer = runtime.n_layer; + materialization.n_expert = runtime.n_expert; + materialization.n_expert_used = n_expert_used; + materialization.primary_placement = placements[0].placement; + materialization.targets.reserve(runtime.targets.size()); + + int explicit_non_cpu_targets = 0; + for (size_t i = 0; i < runtime.targets.size(); ++i) { + ExpertSplitMaterializationTarget target; + target.target_index = (int) i; + target.target = runtime.targets[i].target; + target.placement = std::move(placements[i].placement); + if (target.target.backend != "cpu") { + ++explicit_non_cpu_targets; + } + materialization.targets.push_back(std::move(target)); + } + + materialization.ordered_cold_union = + explicit_non_cpu_targets > 1 && materialization.targets.size() > 1; + if (!materialization.ordered_cold_union) { + out = std::move(materialization); + return true; + } + + materialization.cold_expert_ids_by_layer.resize((size_t) runtime.n_layer); + for (int layer = 0; layer < runtime.n_layer; ++layer) { + std::vector seen((size_t) runtime.n_expert, 0); + const auto & primary_hot = + materialization.primary_placement.hot_expert_ids[(size_t) layer]; + for (int32_t expert : primary_hot) { + if (expert < 0 || expert >= runtime.n_expert) { + if (err) *err = "primary placement expert id out of range"; + return false; + } + seen[(size_t) expert] = 1; + } + + auto & cold = materialization.cold_expert_ids_by_layer[(size_t) layer]; + cold.reserve((size_t) runtime.n_expert - primary_hot.size()); + for (size_t ti = 1; ti < runtime.targets.size(); ++ti) { + const ExpertSplitLayerTarget * layer_target = + runtime.layer_target_ptr((int) ti, layer); + if (!layer_target) { + if (err) *err = "expert split runtime layer target lookup failed"; + return false; + } + for (int32_t expert : layer_target->global_expert_ids) { + if (expert < 0 || expert >= runtime.n_expert) { + if (err) *err = "materialized cold expert id out of range"; + return false; + } + if (seen[(size_t) expert]) { + if (err) { + std::ostringstream ss; + ss << "materialized cold expert duplicated on layer " + << layer << " expert " << expert; + *err = ss.str(); + } + return false; + } + seen[(size_t) expert] = 1; + cold.push_back(expert); + } + } + + const size_t expected = + (size_t) runtime.n_expert - primary_hot.size(); + if (cold.size() != expected) { + if (err) { + std::ostringstream ss; + ss << "materialized cold union count mismatch on layer " + << layer << " expected=" << expected + << " actual=" << cold.size(); + *err = ss.str(); + } + return false; + } + if (std::find(seen.begin(), seen.end(), 0) != seen.end()) { + if (err) { + std::ostringstream ss; + ss << "materialized cold union missing experts on layer " + << layer; + *err = ss.str(); + } + return false; + } + } + + out = std::move(materialization); + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/expert_split_materialization.h b/server/src/common/expert_split_materialization.h new file mode 100644 index 000000000..22e3e5910 --- /dev/null +++ b/server/src/common/expert_split_materialization.h @@ -0,0 +1,40 @@ +// Planner-to-storage bridge for ordered multi-target MoE placements. +// +// The planner/runtime layers decide which target owns each routed expert. This +// materialization step derives the primary local placement plus the ordered +// union of non-primary experts so backend-local storage can preserve the same +// target order before handing work to remote IPC targets or CPU fallback. + +#pragma once + +#include "expert_split_runtime.h" + +#include +#include + +namespace dflash::common { + +struct ExpertSplitMaterializationTarget { + int target_index = -1; + ExpertSplitTarget target; + MoeHybridPlacement placement; +}; + +struct ExpertSplitMaterialization { + int n_layer = 0; + int n_expert = 0; + int n_expert_used = 0; + MoeHybridPlacement primary_placement; + std::vector> cold_expert_ids_by_layer; + std::vector targets; + bool ordered_cold_union = false; + + bool matches(int n_layer_, int n_expert_, int n_expert_used_) const; +}; + +bool build_expert_split_materialization(const ExpertSplitRuntime & runtime, + int n_expert_used, + ExpertSplitMaterialization & out, + std::string * err = nullptr); + +} // namespace dflash::common diff --git a/server/src/common/expert_split_plan.cpp b/server/src/common/expert_split_plan.cpp new file mode 100644 index 000000000..8c091b363 --- /dev/null +++ b/server/src/common/expert_split_plan.cpp @@ -0,0 +1,307 @@ +#include "expert_split_plan.h" + +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +struct RankedUnitRef { + size_t unit_index = 0; + double value = 0.0; +}; + +uint64_t usable_bytes(const ExpertSplitTarget & target) { + if (target.unlimited) { + return std::numeric_limits::max(); + } + if (target.capacity_bytes <= target.reserved_bytes) { + return 0; + } + return target.capacity_bytes - target.reserved_bytes; +} + +bool better_ranked_unit(const ExpertSplitUnit & lhs, + double lhs_value, + const ExpertSplitUnit & rhs, + double rhs_value) { + if (lhs_value != rhs_value) return lhs_value > rhs_value; + if (lhs.score != rhs.score) return lhs.score > rhs.score; + if (lhs.bytes != rhs.bytes) return lhs.bytes < rhs.bytes; + if (lhs.layer != rhs.layer) return lhs.layer < rhs.layer; + return lhs.expert < rhs.expert; +} + +bool can_fit(const ExpertSplitTarget & target, + uint64_t used, + uint64_t bytes) { + if (target.unlimited) return true; + const uint64_t avail = usable_bytes(target); + return used <= avail && bytes <= avail - used; +} + +int floor_for_target(const ExpertSplitConfig & cfg, size_t target_index) { + if (target_index >= cfg.min_per_layer_by_target.size()) { + return 0; + } + return std::max(0, cfg.min_per_layer_by_target[target_index]); +} + +void assign_unit(ExpertSplitPlan & plan, + const ExpertSplitUnit & unit, + int target_index) { + ExpertSplitAssignment & dst = plan.at(unit.layer, unit.expert); + dst.target_index = target_index; + dst.bytes = unit.bytes; + dst.score = unit.score; + if (target_index >= 0 && + (size_t) target_index < plan.target_used_bytes.size()) { + plan.target_used_bytes[(size_t) target_index] += unit.bytes; + } +} + +int choose_first_fit_target(const ExpertSplitPlan & plan, + const ExpertSplitUnit & unit) { + for (size_t i = 0; i < plan.targets.size(); ++i) { + if (can_fit(plan.targets[i], plan.target_used_bytes[i], unit.bytes)) { + return (int) i; + } + } + return -1; +} + +bool place_target_floor(ExpertSplitPlan & plan, + const ExpertSplitConfig & cfg, + const std::vector> & units_by_layer, + const std::vector & units, + std::vector & assigned, + size_t target_index, + std::string * err) { + const int floor = floor_for_target(cfg, target_index); + if (floor <= 0) return true; + + for (int layer = 0; layer < cfg.n_layer; ++layer) { + std::vector candidates = units_by_layer[(size_t) layer]; + std::stable_sort(candidates.begin(), candidates.end(), + [&](size_t lhs, size_t rhs) { + const ExpertSplitUnit & lu = units[lhs]; + const ExpertSplitUnit & ru = units[rhs]; + const double lv = lu.bytes == 0 ? 0.0 : lu.score / (double) lu.bytes; + const double rv = ru.bytes == 0 ? 0.0 : ru.score / (double) ru.bytes; + return better_ranked_unit(lu, lv, ru, rv); + }); + + int placed = 0; + for (size_t unit_index : candidates) { + if (placed >= floor) break; + if (assigned[unit_index]) continue; + const ExpertSplitUnit & unit = units[unit_index]; + if (unit.bytes == 0) continue; + if (!can_fit(plan.targets[target_index], + plan.target_used_bytes[target_index], + unit.bytes)) { + continue; + } + assign_unit(plan, unit, (int) target_index); + assigned[unit_index] = 1; + placed++; + } + if (placed < floor) { + if (err) { + std::ostringstream ss; + ss << "target " << target_index << " floor could not fit" + << " for layer " << layer + << " requested=" << floor + << " placed=" << placed; + *err = ss.str(); + } + return false; + } + } + + return true; +} + +bool validate_units(const ExpertSplitConfig & cfg, + const std::vector & units, + std::string * err) { + if (cfg.n_layer <= 0 || cfg.n_expert <= 0) { + if (err) *err = "expert split dimensions must be positive"; + return false; + } + + std::vector seen((size_t) cfg.n_layer * (size_t) cfg.n_expert, 0); + for (const ExpertSplitUnit & unit : units) { + if (unit.layer < 0 || unit.layer >= cfg.n_layer || + unit.expert < 0 || unit.expert >= cfg.n_expert) { + if (err) *err = "expert split unit out of range"; + return false; + } + const size_t idx = + (size_t) unit.layer * (size_t) cfg.n_expert + (size_t) unit.expert; + if (seen[idx]) { + if (err) *err = "duplicate expert split unit"; + return false; + } + seen[idx] = 1; + } + + if (cfg.require_full_grid) { + const size_t expected = (size_t) cfg.n_layer * (size_t) cfg.n_expert; + if (units.size() != expected) { + if (err) *err = "expert split units do not cover full layer/expert grid"; + return false; + } + if (std::find(seen.begin(), seen.end(), 0) != seen.end()) { + if (err) *err = "expert split units missing from full layer/expert grid"; + return false; + } + } + + return true; +} + +bool validate_targets(const std::vector & targets, + std::string * err) { + for (const ExpertSplitTarget & target : targets) { + if (!target.unlimited && target.capacity_bytes < target.reserved_bytes) { + if (err) *err = "expert split target reserved bytes exceed capacity"; + return false; + } + } + return true; +} + +} // namespace + +uint64_t ExpertSplitTarget::usable_bytes() const { + return dflash::common::usable_bytes(*this); +} + +bool ExpertSplitPlan::matches(int n_layer_, int n_expert_) const { + return n_layer == n_layer_ && n_expert == n_expert_ && + (int) assignments.size() == n_layer * n_expert; +} + +int ExpertSplitPlan::index(int layer, int expert) const { + if (layer < 0 || layer >= n_layer || expert < 0 || expert >= n_expert) { + return -1; + } + return layer * n_expert + expert; +} + +const ExpertSplitAssignment & ExpertSplitPlan::at(int layer, int expert) const { + return assignments[(size_t) index(layer, expert)]; +} + +ExpertSplitAssignment & ExpertSplitPlan::at(int layer, int expert) { + return assignments[(size_t) index(layer, expert)]; +} + +int ExpertSplitPlan::count_on_target(int target_index) const { + int out = 0; + for (const ExpertSplitAssignment & assignment : assignments) { + if (assignment.target_index == target_index) out++; + } + return out; +} + +int ExpertSplitPlan::layer_count_on_target(int layer, int target_index) const { + if (layer < 0 || layer >= n_layer) return 0; + int out = 0; + for (int expert = 0; expert < n_expert; ++expert) { + if (at(layer, expert).target_index == target_index) out++; + } + return out; +} + +bool build_expert_split_plan(const ExpertSplitConfig & cfg, + const std::vector & targets, + const std::vector & units, + ExpertSplitPlan & out, + std::string * err) { + if (!validate_units(cfg, units, err) || + !validate_targets(targets, err)) { + return false; + } + + std::vector plan_targets = targets; + const bool has_cpu_fallback = std::any_of(plan_targets.begin(), plan_targets.end(), + [](const ExpertSplitTarget & target) { + return target.backend == "cpu"; + }); + if (!has_cpu_fallback && cfg.allow_implicit_cpu_fallback) { + plan_targets.push_back({"cpu", "cpu", -1, 0, 0, true}); + } + if (plan_targets.empty()) { + if (err) *err = "expert split requires at least one target"; + return false; + } + + ExpertSplitPlan plan; + plan.n_layer = cfg.n_layer; + plan.n_expert = cfg.n_expert; + plan.targets = std::move(plan_targets); + plan.target_used_bytes.assign(plan.targets.size(), 0); + plan.assignments.assign((size_t) cfg.n_layer * (size_t) cfg.n_expert, {}); + + std::vector sorted_units = units; + std::stable_sort(sorted_units.begin(), sorted_units.end(), + [](const ExpertSplitUnit & lhs, const ExpertSplitUnit & rhs) { + if (lhs.layer != rhs.layer) return lhs.layer < rhs.layer; + return lhs.expert < rhs.expert; + }); + + std::vector> units_by_layer((size_t) cfg.n_layer); + for (size_t i = 0; i < sorted_units.size(); ++i) { + units_by_layer[(size_t) sorted_units[i].layer].push_back(i); + } + std::vector assigned(sorted_units.size(), 0); + + for (size_t target_index = 0; target_index < plan.targets.size(); ++target_index) { + if (!place_target_floor(plan, cfg, units_by_layer, sorted_units, assigned, + target_index, err)) { + return false; + } + } + + std::vector ranked; + ranked.reserve(sorted_units.size()); + for (size_t i = 0; i < sorted_units.size(); ++i) { + if (assigned[i]) continue; + const ExpertSplitUnit & unit = sorted_units[i]; + if (unit.bytes == 0) continue; + ranked.push_back({i, unit.score / (double) unit.bytes}); + } + std::stable_sort(ranked.begin(), ranked.end(), + [&](const RankedUnitRef & lhs, const RankedUnitRef & rhs) { + return better_ranked_unit(sorted_units[lhs.unit_index], lhs.value, + sorted_units[rhs.unit_index], rhs.value); + }); + + for (const RankedUnitRef & ref : ranked) { + if (assigned[ref.unit_index]) continue; + const ExpertSplitUnit & unit = sorted_units[ref.unit_index]; + const int target_index = choose_first_fit_target(plan, unit); + if (target_index < 0) { + if (err) { + std::ostringstream ss; + ss << "expert split capacity exhausted while placing layer " + << unit.layer << " expert " << unit.expert + << " bytes=" << unit.bytes; + *err = ss.str(); + } + return false; + } + assign_unit(plan, unit, target_index); + assigned[ref.unit_index] = 1; + } + + out = std::move(plan); + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/expert_split_plan.h b/server/src/common/expert_split_plan.h new file mode 100644 index 000000000..52a8b0cac --- /dev/null +++ b/server/src/common/expert_split_plan.h @@ -0,0 +1,76 @@ +// Generic ordered expert-residency planner for sparse MoE models. +// +// This planner is intentionally model-agnostic. Adapters describe routed expert +// units and an ordered list of expert targets (for example cuda:0, cuda:1, +// hip:0, cpu), then consume the resulting placement plan. Target order is the +// user-facing tier definition: earlier targets receive hotter experts first. + +#pragma once + +#include +#include +#include + +namespace dflash::common { + +struct ExpertSplitTarget { + std::string name; // Logical name, e.g. "cuda:0" or "cpu". + std::string backend; // "cuda", "hip", "cpu", ... + int device_id = -1; // Backend-local device id, or -1 for host/implicit. + uint64_t capacity_bytes = 0; + uint64_t reserved_bytes = 0; + bool unlimited = false; + + uint64_t usable_bytes() const; +}; + +struct ExpertSplitUnit { + int layer = 0; + int expert = 0; + uint64_t bytes = 0; + double score = 0.0; // Routing count, calibrated probability, or hotness. +}; + +struct ExpertSplitConfig { + int n_layer = 0; + int n_expert = 0; + + // Optional per-target floor. Entry i means target i should receive at + // least this many distinct experts from each layer before later targets + // get considered. Missing entries default to 0. + std::vector min_per_layer_by_target; + + bool allow_implicit_cpu_fallback = true; + bool require_full_grid = true; +}; + +struct ExpertSplitAssignment { + int target_index = -1; // Index into ExpertSplitPlan::targets. + uint64_t bytes = 0; + double score = 0.0; + + bool assigned() const { return target_index >= 0; } +}; + +struct ExpertSplitPlan { + int n_layer = 0; + int n_expert = 0; + std::vector targets; + std::vector target_used_bytes; + std::vector assignments; + + bool matches(int n_layer_, int n_expert_) const; + int index(int layer, int expert) const; + const ExpertSplitAssignment & at(int layer, int expert) const; + ExpertSplitAssignment & at(int layer, int expert); + int count_on_target(int target_index) const; + int layer_count_on_target(int layer, int target_index) const; +}; + +bool build_expert_split_plan(const ExpertSplitConfig & cfg, + const std::vector & targets, + const std::vector & units, + ExpertSplitPlan & out, + std::string * err = nullptr); + +} // namespace dflash::common diff --git a/server/src/common/expert_split_runtime.cpp b/server/src/common/expert_split_runtime.cpp new file mode 100644 index 000000000..80a7fa6b2 --- /dev/null +++ b/server/src/common/expert_split_runtime.cpp @@ -0,0 +1,279 @@ +#include "expert_split_runtime.h" + +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +bool better_local_order(const ExpertSplitAssignment & lhs, + int lhs_expert, + const ExpertSplitAssignment & rhs, + int rhs_expert) { + if (lhs.score != rhs.score) return lhs.score > rhs.score; + if (lhs.bytes != rhs.bytes) return lhs.bytes < rhs.bytes; + return lhs_expert < rhs_expert; +} + +bool validate_plan(const ExpertSplitPlan & plan, std::string * err) { + if (!plan.matches(plan.n_layer, plan.n_expert) || + plan.n_layer <= 0 || plan.n_expert <= 0) { + if (err) *err = "expert split plan not initialized"; + return false; + } + if (plan.targets.empty()) { + if (err) *err = "expert split plan has no targets"; + return false; + } + if (plan.target_used_bytes.size() != plan.targets.size()) { + if (err) *err = "expert split plan target byte accounting mismatch"; + return false; + } + + std::vector recomputed(plan.targets.size(), 0); + for (int layer = 0; layer < plan.n_layer; ++layer) { + for (int expert = 0; expert < plan.n_expert; ++expert) { + const ExpertSplitAssignment & assignment = plan.at(layer, expert); + if (!assignment.assigned()) { + if (err) { + std::ostringstream ss; + ss << "expert split plan missing assignment for layer " + << layer << " expert " << expert; + *err = ss.str(); + } + return false; + } + if (assignment.target_index < 0 || + (size_t) assignment.target_index >= plan.targets.size()) { + if (err) { + std::ostringstream ss; + ss << "expert split plan target index out of range for layer " + << layer << " expert " << expert; + *err = ss.str(); + } + return false; + } + recomputed[(size_t) assignment.target_index] += assignment.bytes; + } + } + + for (size_t i = 0; i < plan.targets.size(); ++i) { + if (recomputed[i] != plan.target_used_bytes[i]) { + if (err) { + std::ostringstream ss; + ss << "expert split plan target byte accounting drift on target " + << i << " expected=" << plan.target_used_bytes[i] + << " actual=" << recomputed[i]; + *err = ss.str(); + } + return false; + } + } + + return true; +} + +} // namespace + +bool ExpertSplitRuntime::matches(int n_layer_, int n_expert_) const { + return n_layer == n_layer_ && + n_expert == n_expert_ && + (int) target_index_by_global.size() == n_layer * n_expert && + (int) local_index_by_global.size() == n_layer * n_expert; +} + +int ExpertSplitRuntime::index(int layer, int expert) const { + if (layer < 0 || layer >= n_layer || expert < 0 || expert >= n_expert) { + return -1; + } + return layer * n_expert + expert; +} + +int ExpertSplitRuntime::target_index(int layer, int expert) const { + const int idx = index(layer, expert); + return idx >= 0 ? target_index_by_global[(size_t) idx] : -1; +} + +int ExpertSplitRuntime::local_index(int layer, int expert) const { + const int idx = index(layer, expert); + return idx >= 0 ? local_index_by_global[(size_t) idx] : -1; +} + +const ExpertSplitRuntimeTarget * ExpertSplitRuntime::target_ptr(int target_index_) const { + if (target_index_ < 0 || (size_t) target_index_ >= targets.size()) { + return nullptr; + } + return &targets[(size_t) target_index_]; +} + +const ExpertSplitLayerTarget * ExpertSplitRuntime::layer_target_ptr(int target_index_, + int layer) const { + const ExpertSplitRuntimeTarget * target = target_ptr(target_index_); + if (!target || layer < 0 || layer >= n_layer) { + return nullptr; + } + return &target->layers[(size_t) layer]; +} + +bool build_expert_split_runtime(const ExpertSplitPlan & plan, + ExpertSplitRuntime & out, + std::string * err) { + if (!validate_plan(plan, err)) { + return false; + } + + ExpertSplitRuntime runtime; + runtime.n_layer = plan.n_layer; + runtime.n_expert = plan.n_expert; + runtime.targets.resize(plan.targets.size()); + runtime.target_index_by_global.assign( + (size_t) runtime.n_layer * (size_t) runtime.n_expert, -1); + runtime.local_index_by_global.assign( + (size_t) runtime.n_layer * (size_t) runtime.n_expert, -1); + + for (size_t target_index = 0; target_index < plan.targets.size(); ++target_index) { + ExpertSplitRuntimeTarget & dst = runtime.targets[target_index]; + dst.target = plan.targets[target_index]; + dst.used_bytes = plan.target_used_bytes[target_index]; + dst.layers.resize((size_t) runtime.n_layer); + for (int layer = 0; layer < runtime.n_layer; ++layer) { + ExpertSplitLayerTarget & layer_dst = dst.layers[(size_t) layer]; + layer_dst.layer = layer; + layer_dst.target_index = (int) target_index; + layer_dst.local_by_global.assign((size_t) runtime.n_expert, -1); + } + } + + for (int layer = 0; layer < runtime.n_layer; ++layer) { + for (int expert = 0; expert < runtime.n_expert; ++expert) { + const ExpertSplitAssignment & assignment = plan.at(layer, expert); + ExpertSplitRuntimeTarget & target = + runtime.targets[(size_t) assignment.target_index]; + ExpertSplitLayerTarget & layer_target = + target.layers[(size_t) layer]; + layer_target.global_expert_ids.push_back((int32_t) expert); + layer_target.total_bytes += assignment.bytes; + layer_target.total_score += assignment.score; + target.total_experts++; + target.total_score += assignment.score; + } + } + + for (size_t target_index = 0; target_index < runtime.targets.size(); ++target_index) { + ExpertSplitRuntimeTarget & target = runtime.targets[target_index]; + for (int layer = 0; layer < runtime.n_layer; ++layer) { + ExpertSplitLayerTarget & layer_target = target.layers[(size_t) layer]; + std::stable_sort(layer_target.global_expert_ids.begin(), + layer_target.global_expert_ids.end(), + [&](int32_t lhs, int32_t rhs) { + return better_local_order(plan.at(layer, lhs), lhs, + plan.at(layer, rhs), rhs); + }); + for (size_t local = 0; local < layer_target.global_expert_ids.size(); ++local) { + const int32_t expert = layer_target.global_expert_ids[local]; + layer_target.local_by_global[(size_t) expert] = (int32_t) local; + const int flat = runtime.index(layer, expert); + runtime.target_index_by_global[(size_t) flat] = (int32_t) target_index; + runtime.local_index_by_global[(size_t) flat] = (int32_t) local; + } + } + } + + out = std::move(runtime); + return true; +} + +bool build_expert_split_target_placement(const ExpertSplitRuntime & runtime, + int n_expert_used, + int target_index_, + MoeHybridPlacement & out, + std::string * err) { + if (!runtime.matches(runtime.n_layer, runtime.n_expert) || + runtime.n_layer <= 0 || runtime.n_expert <= 0) { + if (err) *err = "expert split runtime not initialized"; + return false; + } + if (n_expert_used <= 0 || n_expert_used > runtime.n_expert) { + if (err) *err = "invalid n_expert_used for expert split target placement"; + return false; + } + const ExpertSplitRuntimeTarget * target = runtime.target_ptr(target_index_); + if (!target) { + if (err) *err = "expert split runtime target index out of range"; + return false; + } + if ((int) target->layers.size() != runtime.n_layer) { + if (err) *err = "expert split runtime target layer count mismatch"; + return false; + } + + MoeHybridPlacement placement; + placement.n_layer = runtime.n_layer; + placement.n_expert = runtime.n_expert; + placement.n_expert_used = n_expert_used; + placement.total_hot = target->total_experts; + placement.hot_counts.assign((size_t) runtime.n_layer, 0); + placement.hot_expert_ids.resize((size_t) runtime.n_layer); + + for (int layer = 0; layer < runtime.n_layer; ++layer) { + const ExpertSplitLayerTarget & layer_target = target->layers[(size_t) layer]; + if (layer_target.layer != layer || layer_target.target_index != target_index_) { + if (err) { + std::ostringstream ss; + ss << "expert split runtime layer target mismatch target=" + << target_index_ << " layer=" << layer; + *err = ss.str(); + } + return false; + } + if ((int) layer_target.local_by_global.size() != runtime.n_expert) { + if (err) *err = "expert split runtime local_by_global size mismatch"; + return false; + } + + placement.hot_counts[(size_t) layer] = (int) layer_target.global_expert_ids.size(); + auto & hot = placement.hot_expert_ids[(size_t) layer]; + hot.reserve(layer_target.global_expert_ids.size()); + for (int32_t expert : layer_target.global_expert_ids) { + if (expert < 0 || expert >= runtime.n_expert) { + if (err) *err = "expert split runtime global expert id out of range"; + return false; + } + hot.push_back(expert); + } + } + + out = std::move(placement); + return true; +} + +bool build_all_expert_split_target_placements(const ExpertSplitRuntime & runtime, + int n_expert_used, + std::vector & out, + std::string * err) { + if (!runtime.matches(runtime.n_layer, runtime.n_expert)) { + if (err) *err = "expert split runtime not initialized"; + return false; + } + + std::vector placements; + placements.reserve(runtime.targets.size()); + for (size_t target_index = 0; target_index < runtime.targets.size(); ++target_index) { + ExpertSplitTargetPlacement placement; + placement.target_index = (int) target_index; + if (!build_expert_split_target_placement(runtime, n_expert_used, + placement.target_index, + placement.placement, err)) { + return false; + } + placements.push_back(std::move(placement)); + } + + out = std::move(placements); + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/expert_split_runtime.h b/server/src/common/expert_split_runtime.h new file mode 100644 index 000000000..44715f2d8 --- /dev/null +++ b/server/src/common/expert_split_runtime.h @@ -0,0 +1,72 @@ +// Materialized per-target expert layout derived from ExpertSplitPlan. +// +// This is the bridge between model-agnostic placement planning and later +// backend-specific storage/compute runtimes. Target order stays identical to +// the planner output, and each target receives per-layer local expert indices. + +#pragma once + +#include "expert_split_plan.h" +#include "moe_hybrid_placement.h" + +#include +#include +#include + +namespace dflash::common { + +struct ExpertSplitLayerTarget { + int layer = -1; + int target_index = -1; + std::vector global_expert_ids; + std::vector local_by_global; + uint64_t total_bytes = 0; + double total_score = 0.0; + + bool empty() const { return global_expert_ids.empty(); } +}; + +struct ExpertSplitRuntimeTarget { + ExpertSplitTarget target; + std::vector layers; + int total_experts = 0; + uint64_t used_bytes = 0; + double total_score = 0.0; +}; + +struct ExpertSplitTargetPlacement { + int target_index = -1; + MoeHybridPlacement placement; +}; + +struct ExpertSplitRuntime { + int n_layer = 0; + int n_expert = 0; + std::vector targets; + std::vector target_index_by_global; + std::vector local_index_by_global; + + bool matches(int n_layer_, int n_expert_) const; + int index(int layer, int expert) const; + int target_index(int layer, int expert) const; + int local_index(int layer, int expert) const; + const ExpertSplitRuntimeTarget * target_ptr(int target_index) const; + const ExpertSplitLayerTarget * layer_target_ptr(int target_index, int layer) const; +}; + +bool build_expert_split_runtime(const ExpertSplitPlan & plan, + ExpertSplitRuntime & out, + std::string * err = nullptr); + +bool build_expert_split_target_placement(const ExpertSplitRuntime & runtime, + int n_expert_used, + int target_index, + MoeHybridPlacement & out, + std::string * err = nullptr); + +bool build_all_expert_split_target_placements(const ExpertSplitRuntime & runtime, + int n_expert_used, + std::vector & out, + std::string * err = nullptr); + +} // namespace dflash::common diff --git a/server/src/common/expert_split_state.cpp b/server/src/common/expert_split_state.cpp new file mode 100644 index 000000000..6e1988c1e --- /dev/null +++ b/server/src/common/expert_split_state.cpp @@ -0,0 +1,124 @@ +#include "expert_split_state.h" + +namespace dflash::common { + +bool ExpertSplitStateComponents::empty() const { + return plan.assignments.empty() && + runtime.targets.empty() && + compute_runtime.targets.empty() && + materialization.targets.empty(); +} + +bool ExpertSplitStateComponents::matches(int n_layer_, + int n_expert_, + int n_expert_used_) const { + return plan.matches(n_layer_, n_expert_) && + runtime.matches(n_layer_, n_expert_) && + compute_runtime.matches(n_layer_, n_expert_, n_expert_used_) && + materialization.matches(n_layer_, n_expert_, n_expert_used_); +} + +int ExpertSplitLayerMapping::split_layer_for_physical(int physical_layer) const { + if (physical_layer < 0 || + (size_t) physical_layer >= split_layer_by_physical_layer.size()) { + return -1; + } + return split_layer_by_physical_layer[(size_t) physical_layer]; +} + +bool build_expert_split_state(const ExpertSplitConfig & cfg, + const std::vector & targets, + const std::vector & units, + int n_expert_used, + ExpertSplitStateComponents & out, + std::string * err) { + out = ExpertSplitStateComponents{}; + + if (!build_expert_split_plan(cfg, targets, units, out.plan, err)) { + out = ExpertSplitStateComponents{}; + return false; + } + if (!build_expert_split_runtime(out.plan, out.runtime, err)) { + out = ExpertSplitStateComponents{}; + return false; + } + if (!build_expert_split_compute_runtime( + out.runtime, n_expert_used, out.compute_runtime, err)) { + out = ExpertSplitStateComponents{}; + return false; + } + if (!build_expert_split_materialization( + out.runtime, n_expert_used, out.materialization, err)) { + out = ExpertSplitStateComponents{}; + return false; + } + return true; +} + +bool build_expert_split_layer_mapping( + int n_total_layer, + const std::vector & physical_layer_by_split_layer, + ExpertSplitLayerMapping & out, + std::string * err) { + out = ExpertSplitLayerMapping{}; + if (n_total_layer <= 0) { + if (err) *err = "expert split layer mapping requires n_total_layer > 0"; + return false; + } + + out.n_total_layer = n_total_layer; + out.physical_layer_by_split_layer = physical_layer_by_split_layer; + out.split_layer_by_physical_layer.assign((size_t) n_total_layer, -1); + + for (size_t split_layer = 0; + split_layer < physical_layer_by_split_layer.size(); + ++split_layer) { + const int32_t physical_layer = + physical_layer_by_split_layer[split_layer]; + if (physical_layer < 0 || physical_layer >= n_total_layer) { + out = ExpertSplitLayerMapping{}; + if (err) *err = "expert split physical layer out of range"; + return false; + } + const int32_t existing = + out.split_layer_by_physical_layer[(size_t) physical_layer]; + if (existing >= 0) { + out = ExpertSplitLayerMapping{}; + if (err) *err = "expert split physical layer mapped more than once"; + return false; + } + out.split_layer_by_physical_layer[(size_t) physical_layer] = + (int32_t) split_layer; + } + return true; +} + +bool build_contiguous_expert_split_layer_mapping( + int n_total_layer, + int first_split_layer, + int n_split_layer, + ExpertSplitLayerMapping & out, + std::string * err) { + if (n_total_layer <= 0) { + if (err) *err = "expert split layer mapping requires n_total_layer > 0"; + return false; + } + if (first_split_layer < 0 || first_split_layer > n_total_layer) { + if (err) *err = "expert split first_split_layer out of range"; + return false; + } + if (n_split_layer < 0 || n_split_layer > n_total_layer - first_split_layer) { + if (err) *err = "expert split contiguous layer count out of range"; + return false; + } + + std::vector physical_layer_by_split_layer((size_t) n_split_layer); + for (int split_layer = 0; split_layer < n_split_layer; ++split_layer) { + physical_layer_by_split_layer[(size_t) split_layer] = + first_split_layer + split_layer; + } + return build_expert_split_layer_mapping( + n_total_layer, physical_layer_by_split_layer, out, err); +} + +} // namespace dflash::common diff --git a/server/src/common/expert_split_state.h b/server/src/common/expert_split_state.h new file mode 100644 index 000000000..52cf217f6 --- /dev/null +++ b/server/src/common/expert_split_state.h @@ -0,0 +1,65 @@ +// Shared expert-split state assembly helpers. +// +// Sparse adapters and backends all build the same four planning/runtime +// products: plan, runtime, compute runtime, and materialized placement views. +// Keep that assembly in one place so model-specific adapters only need to +// provide units, targets, and any layer mapping they require. + +#pragma once + +#include "expert_split_compute_runtime.h" +#include "expert_split_materialization.h" +#include "expert_split_plan.h" +#include "expert_split_runtime.h" + +#include +#include +#include + +namespace dflash::common { + +struct ExpertSplitStateComponents { + ExpertSplitPlan plan; + ExpertSplitRuntime runtime; + ExpertSplitComputeRuntime compute_runtime; + ExpertSplitMaterialization materialization; + + bool empty() const; + bool matches(int n_layer_, int n_expert_, int n_expert_used_) const; +}; + +struct ExpertSplitLayerMapping { + int n_total_layer = 0; + std::vector physical_layer_by_split_layer; + std::vector split_layer_by_physical_layer; + + int split_layer_count() const { + return (int) physical_layer_by_split_layer.size(); + } + int split_layer_for_physical(int physical_layer) const; + bool is_split_layer(int physical_layer) const { + return split_layer_for_physical(physical_layer) >= 0; + } +}; + +bool build_expert_split_state(const ExpertSplitConfig & cfg, + const std::vector & targets, + const std::vector & units, + int n_expert_used, + ExpertSplitStateComponents & out, + std::string * err = nullptr); + +bool build_expert_split_layer_mapping( + int n_total_layer, + const std::vector & physical_layer_by_split_layer, + ExpertSplitLayerMapping & out, + std::string * err = nullptr); + +bool build_contiguous_expert_split_layer_mapping( + int n_total_layer, + int first_split_layer, + int n_split_layer, + ExpertSplitLayerMapping & out, + std::string * err = nullptr); + +} // namespace dflash::common diff --git a/server/src/common/expert_split_target_config.cpp b/server/src/common/expert_split_target_config.cpp new file mode 100644 index 000000000..eed716cc1 --- /dev/null +++ b/server/src/common/expert_split_target_config.cpp @@ -0,0 +1,347 @@ +#include "expert_split_target_config.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +std::string trim_ascii(std::string text) { + size_t begin = 0; + while (begin < text.size() && std::isspace((unsigned char) text[begin])) { + ++begin; + } + size_t end = text.size(); + while (end > begin && std::isspace((unsigned char) text[end - 1])) { + --end; + } + return text.substr(begin, end - begin); +} + +bool parse_u64_text(const std::string & text, uint64_t & out) { + if (text.empty()) return false; + errno = 0; + char * end = nullptr; + unsigned long long value = std::strtoull(text.c_str(), &end, 10); + if (errno == ERANGE || end == text.c_str() || *end != '\0') { + return false; + } + out = (uint64_t) value; + return true; +} + +bool parse_capacity_token(const std::string & raw, uint64_t & out_bytes) { + const std::string text = trim_ascii(raw); + if (text.empty()) return false; + + size_t suffix_begin = text.size(); + while (suffix_begin > 0 && + std::isalpha((unsigned char) text[suffix_begin - 1])) { + --suffix_begin; + } + const std::string number = trim_ascii(text.substr(0, suffix_begin)); + const std::string suffix = text.substr(suffix_begin); + uint64_t value = 0; + if (!parse_u64_text(number, value)) { + return false; + } + + uint64_t scale = 1; + if (suffix.empty() || suffix == "b" || suffix == "B") { + scale = 1; + } else if (suffix == "k" || suffix == "K" || suffix == "kb" || suffix == "KB") { + scale = 1024ULL; + } else if (suffix == "m" || suffix == "M" || suffix == "mb" || suffix == "MB") { + scale = 1024ULL * 1024ULL; + } else if (suffix == "g" || suffix == "G" || suffix == "gb" || suffix == "GB") { + scale = 1024ULL * 1024ULL * 1024ULL; + } else { + return false; + } + if (value > std::numeric_limits::max() / scale) { + return false; + } + out_bytes = value * scale; + return true; +} + +bool query_backend_device_total_memory(PlacementBackend backend, + int device_id, + uint64_t & total_bytes) { + if (device_id < 0) return false; + const char * reg_name = nullptr; + switch (backend) { + case PlacementBackend::Auto: + reg_name = placement_backend_name(compiled_placement_backend()); + break; + case PlacementBackend::Cuda: + case PlacementBackend::Hip: + reg_name = placement_backend_name(backend); + break; + } + if (!reg_name) return false; + + ggml_backend_reg_t reg = ggml_backend_reg_by_name(reg_name); + if (!reg) return false; + const size_t n_dev = ggml_backend_reg_dev_count(reg); + if (device_id < 0 || (size_t)device_id >= n_dev) return false; + ggml_backend_dev_t dev = ggml_backend_reg_dev_get(reg, (size_t)device_id); + if (!dev) return false; + size_t free_bytes = 0; + size_t total = 0; + ggml_backend_dev_memory(dev, &free_bytes, &total); + (void) free_bytes; + if (total == 0) return false; + total_bytes = (uint64_t) total; + return true; +} + +bool specs_have_duplicate_device(const std::vector & specs) { + for (size_t i = 0; i < specs.size(); ++i) { + for (size_t j = i + 1; j < specs.size(); ++j) { + if (specs[i].backend == specs[j].backend && + specs[i].device_id == specs[j].device_id) { + return true; + } + } + } + return false; +} + +} // namespace + +bool parse_expert_split_target_list(const std::string & value, + std::vector & out, + std::string * err) { + out.clear(); + if (value.empty()) { + if (err) *err = "expert target list is empty"; + return false; + } + + size_t begin = 0; + while (begin < value.size()) { + const size_t end = value.find(',', begin); + const std::string item = trim_ascii(value.substr( + begin, end == std::string::npos ? std::string::npos : end - begin)); + if (item.empty()) { + if (err) *err = "expert target list contains an empty item"; + out.clear(); + return false; + } + + ExpertSplitTargetSpec spec; + if (item == "cpu") { + spec.backend = PlacementBackend::Auto; + spec.device_id = -1; + spec.auto_capacity = false; + spec.capacity_bytes = 0; + spec.unlimited = true; + } else { + const size_t sep = item.find(':'); + if (sep == std::string::npos || sep == 0 || sep + 1 >= item.size()) { + if (err) *err = "expert target item must look like backend:id or cpu"; + out.clear(); + return false; + } + if (!parse_placement_backend(item.substr(0, sep), spec.backend)) { + if (err) *err = "unsupported expert target backend: " + item.substr(0, sep); + out.clear(); + return false; + } + char * num_end = nullptr; + long device_id = std::strtol(item.c_str() + sep + 1, &num_end, 10); + if (num_end == item.c_str() + sep + 1 || *num_end != '\0' || device_id < 0) { + if (err) *err = "invalid expert target device id: " + item; + out.clear(); + return false; + } + spec.device_id = (int)device_id; + spec.auto_capacity = true; + spec.unlimited = false; + } + + out.push_back(spec); + if (end == std::string::npos) break; + begin = end + 1; + } + + if (out.empty()) { + if (err) *err = "expert target list is empty"; + return false; + } + if (specs_have_duplicate_device(out)) { + if (err) *err = "expert target list contains duplicate devices"; + out.clear(); + return false; + } + return true; +} + +bool parse_expert_split_capacity_overrides(const std::string & value, + std::vector & out_bytes, + std::string * err) { + out_bytes.clear(); + if (value.empty()) return true; + + size_t begin = 0; + while (begin < value.size()) { + const size_t end = value.find(',', begin); + const std::string item = trim_ascii(value.substr( + begin, end == std::string::npos ? std::string::npos : end - begin)); + if (item.empty()) { + if (err) *err = "expert target capacity list contains an empty item"; + out_bytes.clear(); + return false; + } + uint64_t bytes = 0; + if (item == "auto" || item == "AUTO") { + out_bytes.push_back(0); + } else if (!parse_capacity_token(item, bytes)) { + if (err) *err = "invalid expert target capacity: " + item; + out_bytes.clear(); + return false; + } else { + out_bytes.push_back(bytes); + } + if (end == std::string::npos) break; + begin = end + 1; + } + + return true; +} + +bool build_expert_split_targets(const std::vector & specs, + uint64_t primary_capacity_bytes, + std::vector & out, + std::string * err) { + out.clear(); + if (specs.empty()) { + if (err) *err = "no expert targets configured"; + return false; + } + + out.reserve(specs.size()); + for (size_t i = 0; i < specs.size(); ++i) { + const ExpertSplitTargetSpec & spec = specs[i]; + ExpertSplitTarget target; + if (spec.unlimited) { + target.name = "cpu"; + target.backend = "cpu"; + target.device_id = -1; + target.unlimited = true; + out.push_back(std::move(target)); + continue; + } + + PlacementBackend backend = + spec.backend == PlacementBackend::Auto + ? compiled_placement_backend() + : spec.backend; + target.name = std::string(placement_backend_name(backend)) + ":" + + std::to_string(spec.device_id); + target.backend = placement_backend_name(backend); + target.device_id = spec.device_id; + target.unlimited = false; + + uint64_t capacity_bytes = 0; + if (i == 0 && primary_capacity_bytes > 0) { + capacity_bytes = primary_capacity_bytes; + } else if (!spec.auto_capacity) { + capacity_bytes = spec.capacity_bytes; + } else if (spec.capacity_bytes > 0) { + capacity_bytes = spec.capacity_bytes; + } else if (!query_backend_device_total_memory(backend, spec.device_id, capacity_bytes)) { + if (err) { + *err = "could not auto-discover capacity for expert target " + + target.name; + } + out.clear(); + return false; + } + + target.capacity_bytes = capacity_bytes; + out.push_back(std::move(target)); + } + + return true; +} + +bool resolve_expert_split_targets_from_env(const char * targets_env_name, + const char * caps_env_name, + uint64_t primary_capacity_bytes, + std::vector & out, + std::string * err) { + out.clear(); + if (!targets_env_name || !*targets_env_name) { + if (err) *err = "expert target env name is empty"; + return false; + } + + const char * targets_env = std::getenv(targets_env_name); + if (!targets_env || !*targets_env) { + return true; + } + + std::vector specs; + if (!parse_expert_split_target_list(targets_env, specs, err)) { + return false; + } + + if (caps_env_name && *caps_env_name) { + if (const char * caps_env = std::getenv(caps_env_name)) { + std::vector caps; + if (!parse_expert_split_capacity_overrides(caps_env, caps, err)) { + return false; + } + if (!caps.empty() && caps.size() != specs.size()) { + if (err) { + *err = "expert target cap count must match expert target count"; + } + return false; + } + for (size_t i = 0; i < caps.size(); ++i) { + if (caps[i] == 0) continue; + specs[i].auto_capacity = false; + specs[i].capacity_bytes = caps[i]; + } + } + } + + return build_expert_split_targets(specs, primary_capacity_bytes, out, err); +} + +bool validate_primary_expert_split_target( + const std::vector & targets, + PlacementBackend local_backend, + int local_device_id, + std::string * err) { + if (targets.empty()) { + if (err) *err = "expert split requires at least one target"; + return false; + } + const char * backend_name = placement_backend_name(local_backend); + if (!backend_name) { + if (err) *err = "could not resolve local expert backend"; + return false; + } + const ExpertSplitTarget & primary = targets.front(); + if (primary.backend != backend_name || primary.device_id != local_device_id) { + if (err) { + *err = "first expert target must match the local backend device (" + + std::string(backend_name) + ":" + std::to_string(local_device_id) + ")"; + } + return false; + } + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/expert_split_target_config.h b/server/src/common/expert_split_target_config.h new file mode 100644 index 000000000..1e144587d --- /dev/null +++ b/server/src/common/expert_split_target_config.h @@ -0,0 +1,52 @@ +// Ordered expert-target config parsing and target materialization helpers. +// +// This is the planning-side adapter for sparse MoE backends that want an +// ordered target list (for example "cuda:0,cuda:1,cpu") without exposing +// tier-specific naming. It parses user config, optionally auto-discovers +// device capacities, and emits ExpertSplitTarget entries for the planner. + +#pragma once + +#include "expert_split_plan.h" +#include "placement/placement_backend.h" + +#include +#include +#include + +namespace dflash::common { + +struct ExpertSplitTargetSpec { + PlacementBackend backend = PlacementBackend::Auto; + int device_id = -1; + bool auto_capacity = true; + uint64_t capacity_bytes = 0; + bool unlimited = false; +}; + +bool parse_expert_split_target_list(const std::string & value, + std::vector & out, + std::string * err = nullptr); + +bool parse_expert_split_capacity_overrides(const std::string & value, + std::vector & out_bytes, + std::string * err = nullptr); + +bool build_expert_split_targets(const std::vector & specs, + uint64_t primary_capacity_bytes, + std::vector & out, + std::string * err = nullptr); + +bool resolve_expert_split_targets_from_env(const char * targets_env_name, + const char * caps_env_name, + uint64_t primary_capacity_bytes, + std::vector & out, + std::string * err = nullptr); + +bool validate_primary_expert_split_target( + const std::vector & targets, + PlacementBackend local_backend, + int local_device_id, + std::string * err = nullptr); + +} // namespace dflash::common diff --git a/server/src/common/gguf_tensor_data.cpp b/server/src/common/gguf_tensor_data.cpp new file mode 100644 index 000000000..d6f7a3e8f --- /dev/null +++ b/server/src/common/gguf_tensor_data.cpp @@ -0,0 +1,480 @@ +#include "gguf_tensor_data.h" + +#include "common/gguf_mmap.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +uint16_t get_u16_or(const gguf_context * g, const char * key, uint16_t fallback) { + const int64_t id = gguf_find_key(g, key); + if (id < 0) return fallback; + return gguf_get_val_u16(g, id); +} + +bool parse_positive_int(const std::string & text, int & out) { + if (text.empty()) return false; + errno = 0; + char * end = nullptr; + const long v = std::strtol(text.c_str(), &end, 10); + if (errno == ERANGE || end == text.c_str() || *end != '\0' || + v <= 0 || v > std::numeric_limits::max()) { + return false; + } + out = (int)v; + return true; +} + +bool split_prefix_from_path(const std::string & path, + int split_no, + int split_count, + std::string & prefix, + std::string & err) { + static const std::string suffix = ".gguf"; + if (path.size() <= suffix.size() || + path.compare(path.size() - suffix.size(), suffix.size(), suffix) != 0) { + err = "split GGUF path does not end in .gguf: " + path; + return false; + } + + const size_t suffix_pos = path.size() - suffix.size(); + const size_t of_pos = path.rfind("-of-", suffix_pos); + if (of_pos == std::string::npos) { + err = "split GGUF path is missing -of- marker: " + path; + return false; + } + const size_t dash_pos = path.rfind('-', of_pos == 0 ? 0 : of_pos - 1); + if (dash_pos == std::string::npos || dash_pos + 1 >= of_pos) { + err = "split GGUF path is missing split number: " + path; + return false; + } + + int file_no = 0; + int file_count = 0; + if (!parse_positive_int(path.substr(dash_pos + 1, of_pos - dash_pos - 1), file_no) || + !parse_positive_int(path.substr(of_pos + 4, suffix_pos - (of_pos + 4)), file_count)) { + err = "split GGUF path has invalid split numbers: " + path; + return false; + } + if (file_no != split_no + 1 || file_count != split_count) { + std::ostringstream ss; + ss << "split GGUF path metadata mismatch: path says " + << file_no << "-of-" << file_count + << " but metadata says " << (split_no + 1) + << "-of-" << split_count; + err = ss.str(); + return false; + } + + prefix = path.substr(0, dash_pos); + return true; +} + +std::string split_path_from_prefix(const std::string & prefix, + int split_no, + int split_count) { + char tail[64]; + std::snprintf(tail, sizeof(tail), "-%05d-of-%05d.gguf", + split_no + 1, split_count); + return prefix + tail; +} + +} // namespace + +struct GgufTensorDataReader::Impl { + struct Shard { + std::string path; + gguf_context * gctx = nullptr; + ggml_context * tensor_ctx = nullptr; + GgufMmap mmap; + }; + + std::vector> shards; + ggml_context * merged_ctx = nullptr; + + ~Impl() { + if (merged_ctx) { + ggml_free(merged_ctx); + merged_ctx = nullptr; + } + for (auto & shard : shards) { + if (!shard) continue; + if (shard->tensor_ctx) { + ggml_free(shard->tensor_ctx); + shard->tensor_ctx = nullptr; + } + if (shard->gctx) { + gguf_free(shard->gctx); + shard->gctx = nullptr; + } + } + } + + bool open_one(const std::string & path, + bool build_tensor_context, + int expected_no, + int expected_count, + std::string & err) { + auto shard = std::make_unique(); + shard->path = path; + + gguf_init_params gip{}; + gip.no_alloc = true; + if (build_tensor_context) { + gip.ctx = &shard->tensor_ctx; + } + shard->gctx = gguf_init_from_file(path.c_str(), gip); + if (!shard->gctx) { + err = "gguf_init_from_file failed: " + path; + return false; + } + + if (expected_count > 1) { + const uint16_t split_no = + get_u16_or(shard->gctx, "split.no", UINT16_MAX); + const uint16_t split_count = + get_u16_or(shard->gctx, "split.count", 0); + if (split_no != (uint16_t)expected_no || + split_count != (uint16_t)expected_count) { + std::ostringstream ss; + ss << "split metadata mismatch in " << path + << ": split.no=" << split_no + << " split.count=" << split_count + << " expected " << expected_no + << "/" << expected_count; + err = ss.str(); + return false; + } + } + + shards.push_back(std::move(shard)); + return true; + } + + bool build_merged_context(std::string & err) { + size_t n_tensors = 0; + for (const auto & shard : shards) { + n_tensors += (size_t)gguf_get_n_tensors(shard->gctx); + } + if (n_tensors == 0) { + err = "GGUF shard set contains no tensors"; + return false; + } + if (n_tensors > std::numeric_limits::max() / ggml_tensor_overhead()) { + err = "GGUF tensor metadata context size overflow"; + return false; + } + + ggml_init_params ip{}; + ip.mem_size = n_tensors * ggml_tensor_overhead(); + ip.mem_buffer = nullptr; + ip.no_alloc = true; + merged_ctx = ggml_init(ip); + if (!merged_ctx) { + err = "failed to allocate merged GGUF tensor context"; + return false; + } + + for (const auto & shard : shards) { + for (int64_t tid = 0; tid < gguf_get_n_tensors(shard->gctx); ++tid) { + const char * name = gguf_get_tensor_name(shard->gctx, tid); + ggml_tensor * src = shard->tensor_ctx + ? ggml_get_tensor(shard->tensor_ctx, name) + : nullptr; + if (!src) { + err = std::string("missing shard tensor metadata: ") + name; + return false; + } + if (ggml_get_tensor(merged_ctx, name)) { + err = std::string("duplicate tensor across GGUF shards: ") + name; + return false; + } + ggml_tensor * dst = + ggml_new_tensor(merged_ctx, src->type, GGML_MAX_DIMS, src->ne); + if (!dst) { + err = std::string("failed to create merged tensor metadata: ") + name; + return false; + } + ggml_set_name(dst, name); + } + } + return true; + } +}; + +GgufTensorDataReader::GgufTensorDataReader() + : impl_(std::make_unique()) {} + +GgufTensorDataReader::~GgufTensorDataReader() = default; + +bool GgufTensorDataReader::open(const std::string & path, + bool build_merged_tensor_context, + std::string & err) { + impl_ = std::make_unique(); + + gguf_context * probe = gguf_init_from_file(path.c_str(), gguf_init_params{}); + if (!probe) { + err = "gguf_init_from_file failed: " + path; + return false; + } + + const uint16_t split_count = get_u16_or(probe, "split.count", 0); + const uint16_t split_no = get_u16_or(probe, "split.no", 0); + gguf_free(probe); + + if (split_count <= 1) { + if (!impl_->open_one(path, build_merged_tensor_context, + /*expected_no=*/0, /*expected_count=*/1, err)) { + return false; + } + } else { + std::string prefix; + if (!split_prefix_from_path(path, split_no, split_count, prefix, err)) { + return false; + } + for (int i = 0; i < (int)split_count; ++i) { + const std::string shard_path = + split_path_from_prefix(prefix, i, (int)split_count); + if (!impl_->open_one(shard_path, build_merged_tensor_context, + i, (int)split_count, err)) { + return false; + } + } + } + + if (build_merged_tensor_context && + !impl_->build_merged_context(err)) { + return false; + } + return true; +} + +bool GgufTensorDataReader::open_mmaps(std::string & err) { + for (auto & shard : impl_->shards) { + if (!shard->mmap.open(shard->path, err)) { + return false; + } + } + return true; +} + +const gguf_context * GgufTensorDataReader::primary_context() const { + return impl_->shards.empty() ? nullptr : impl_->shards.front()->gctx; +} + +ggml_context * GgufTensorDataReader::merged_context() const { + return impl_->merged_ctx; +} + +ggml_context * GgufTensorDataReader::release_merged_context() { + ggml_context * out = impl_->merged_ctx; + impl_->merged_ctx = nullptr; + return out; +} + +int GgufTensorDataReader::shard_count() const { + return (int)impl_->shards.size(); +} + +const gguf_context * GgufTensorDataReader::shard_context(int shard_index) const { + if (shard_index < 0 || shard_index >= shard_count()) return nullptr; + return impl_->shards[(size_t)shard_index]->gctx; +} + +bool GgufTensorDataReader::find_tensor(const char * name, GgufTensorRef & out) const { + out = {}; + if (!name || !*name) return false; + for (size_t si = 0; si < impl_->shards.size(); ++si) { + const auto * shard = impl_->shards[si].get(); + const int64_t tid = gguf_find_tensor(shard->gctx, name); + if (tid < 0) continue; + + out.shard_index = (int)si; + out.tensor_id = tid; + out.tensor = impl_->merged_ctx ? ggml_get_tensor(impl_->merged_ctx, name) + : nullptr; + out.size = gguf_get_tensor_size(shard->gctx, tid); + out.type = gguf_get_tensor_type(shard->gctx, tid); + if (shard->mmap.is_open()) { + const size_t off = + gguf_get_data_offset(shard->gctx) + + gguf_get_tensor_offset(shard->gctx, tid); + if (off <= shard->mmap.size() && + out.size <= shard->mmap.size() - off) { + out.data = static_cast(shard->mmap.data()) + off; + } + } + return true; + } + return false; +} + +const std::string & GgufTensorDataReader::shard_path(int shard_index) const { + return impl_->shards[(size_t)shard_index]->path; +} + +const void * GgufTensorDataReader::shard_mmap_data(int shard_index) const { + if (shard_index < 0 || shard_index >= shard_count()) return nullptr; + return impl_->shards[(size_t)shard_index]->mmap.data(); +} + +size_t GgufTensorDataReader::shard_mmap_size(int shard_index) const { + if (shard_index < 0 || shard_index >= shard_count()) return 0; + return impl_->shards[(size_t)shard_index]->mmap.size(); +} + +bool GgufTensorDataReader::release_single_mmap( + const void *& data, + size_t & size, + int & fd, + std::string & err) { + data = nullptr; + size = 0; + fd = -1; + if (shard_count() != 1) { + err = "release_single_mmap requires exactly one GGUF shard"; + return false; + } + GgufMmap::OwnedRegion region = impl_->shards.front()->mmap.release(); + data = region.data; + size = region.size; + fd = region.fd; + return data != nullptr && size > 0; +} + +std::vector make_layer_expert_file_data( + const GgufTensorDataReader & reader, + int n_layer) { + std::vector layer_file_data((size_t)n_layer); + for (int il = 0; il < n_layer; ++il) { + char name[128]; + auto find_tensor_data = [&](const char * suffix) -> ExpertTensorFileData { + std::snprintf(name, sizeof(name), "blk.%d.%s.weight", il, suffix); + GgufTensorRef ref; + if (!reader.find_tensor(name, ref) || !ref.data) return {}; + return { ref.data, ref.size }; + }; + layer_file_data[(size_t)il].gate_exps = find_tensor_data("ffn_gate_exps"); + layer_file_data[(size_t)il].up_exps = find_tensor_data("ffn_up_exps"); + layer_file_data[(size_t)il].down_exps = find_tensor_data("ffn_down_exps"); + layer_file_data[(size_t)il].gate_up_exps = find_tensor_data("ffn_gate_up_exps"); + } + return layer_file_data; +} + +int64_t gguf_tensor_ref_row_width(const GgufTensorRef & ref) { + if (!ref.tensor) return 0; + return ref.tensor->ne[0]; +} + +int64_t gguf_tensor_ref_row_count(const GgufTensorRef & ref) { + if (!ref.tensor) return 0; + int64_t rows = 1; + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + rows *= ref.tensor->ne[i]; + } + return rows; +} + +bool gguf_tensor_ref_rows_to_host_f32( + const GgufTensorRef & ref, + int64_t row_begin, + int64_t row_count, + std::vector & out, + std::string * err) { + out.clear(); + if (err) err->clear(); + if (!ref.tensor || ref.tensor_id < 0) { + if (err) *err = "GGUF tensor ref is not initialized"; + return false; + } + if (!ref.data) { + if (err) *err = "GGUF tensor ref has no mapped data"; + return false; + } + + const int64_t row_width = gguf_tensor_ref_row_width(ref); + const int64_t total_rows = gguf_tensor_ref_row_count(ref); + if (row_width <= 0 || total_rows <= 0) { + if (err) *err = "GGUF tensor ref shape is invalid"; + return false; + } + if (row_begin < 0 || row_count < 0 || row_begin > total_rows || + row_count > total_rows - row_begin) { + if (err) *err = "GGUF tensor ref row range is out of bounds"; + return false; + } + if (row_count == 0) return true; + + const ggml_type type = ref.type; + const ggml_type_traits * traits = ggml_get_type_traits(type); + const bool native_f32 = type == GGML_TYPE_F32; + const bool native_f16 = type == GGML_TYPE_F16; + const bool native_bf16 = type == GGML_TYPE_BF16; + if (!(native_f32 || native_f16 || native_bf16) && + (!traits || !traits->to_float)) { + if (err) *err = "GGUF tensor ref type cannot be dequantized to f32"; + return false; + } + + const int64_t block = ggml_blck_size(type); + if (block <= 0 || row_width % block != 0) { + if (err) *err = "GGUF tensor ref row width is incompatible with tensor block size"; + return false; + } + + const size_t row_bytes = ggml_row_size(type, row_width); + if (row_bytes == 0) { + if (err) *err = "GGUF tensor ref row size is zero"; + return false; + } + if ((size_t) total_rows > std::numeric_limits::max() / row_bytes) { + if (err) *err = "GGUF tensor ref row byte size overflow"; + return false; + } + const size_t total_bytes = (size_t) total_rows * row_bytes; + if (total_bytes > ref.size) { + if (err) *err = "GGUF tensor ref byte extent is smaller than tensor rows"; + return false; + } + if ((size_t) row_count > std::numeric_limits::max() / (size_t) row_width) { + if (err) *err = "GGUF tensor ref output size overflow"; + return false; + } + + out.resize((size_t) row_count * (size_t) row_width); + for (int64_t row = 0; row < row_count; ++row) { + const size_t src_off = ((size_t) row_begin + (size_t) row) * row_bytes; + const uint8_t * src = ref.data + src_off; + float * dst = out.data() + (size_t) row * (size_t) row_width; + if (native_f32) { + std::memcpy(dst, src, (size_t) row_width * sizeof(float)); + } else if (native_f16) { + ggml_fp16_to_fp32_row((const ggml_fp16_t *) src, dst, row_width); + } else if (native_bf16) { + ggml_bf16_to_fp32_row((const ggml_bf16_t *) src, dst, row_width); + } else { + traits->to_float(src, dst, row_width); + } + } + return true; +} + +bool gguf_tensor_ref_all_to_host_f32( + const GgufTensorRef & ref, + std::vector & out, + std::string * err) { + return gguf_tensor_ref_rows_to_host_f32( + ref, 0, gguf_tensor_ref_row_count(ref), out, err); +} + +} // namespace dflash::common diff --git a/server/src/common/gguf_tensor_data.h b/server/src/common/gguf_tensor_data.h new file mode 100644 index 000000000..799e8b0f9 --- /dev/null +++ b/server/src/common/gguf_tensor_data.h @@ -0,0 +1,76 @@ +#pragma once + +#include "ggml.h" +#include "gguf.h" +#include "moe_hybrid_storage.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +struct GgufTensorRef { + int shard_index = -1; + int64_t tensor_id = -1; + ggml_tensor * tensor = nullptr; + const uint8_t * data = nullptr; + size_t size = 0; + ggml_type type = GGML_TYPE_COUNT; +}; + +class GgufTensorDataReader { +public: + GgufTensorDataReader(); + ~GgufTensorDataReader(); + + GgufTensorDataReader(const GgufTensorDataReader &) = delete; + GgufTensorDataReader & operator=(const GgufTensorDataReader &) = delete; + + bool open(const std::string & path, + bool build_merged_tensor_context, + std::string & err); + bool open_mmaps(std::string & err); + + const gguf_context * primary_context() const; + ggml_context * merged_context() const; + ggml_context * release_merged_context(); + + int shard_count() const; + const gguf_context * shard_context(int shard_index) const; + + bool find_tensor(const char * name, GgufTensorRef & out) const; + const std::string & shard_path(int shard_index) const; + const void * shard_mmap_data(int shard_index) const; + size_t shard_mmap_size(int shard_index) const; + + // Transfer the sole mmap to the caller. Valid only for single-shard readers. + bool release_single_mmap(const void *& data, size_t & size, int & fd, std::string & err); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +std::vector make_layer_expert_file_data( + const GgufTensorDataReader & reader, + int n_layer); + +int64_t gguf_tensor_ref_row_width(const GgufTensorRef & ref); +int64_t gguf_tensor_ref_row_count(const GgufTensorRef & ref); + +bool gguf_tensor_ref_rows_to_host_f32( + const GgufTensorRef & ref, + int64_t row_begin, + int64_t row_count, + std::vector & out, + std::string * err = nullptr); + +bool gguf_tensor_ref_all_to_host_f32( + const GgufTensorRef & ref, + std::vector & out, + std::string * err = nullptr); + +} // namespace dflash::common diff --git a/server/src/common/moe_expert_compute.cpp b/server/src/common/moe_expert_compute.cpp new file mode 100644 index 000000000..65975948e --- /dev/null +++ b/server/src/common/moe_expert_compute.cpp @@ -0,0 +1,1284 @@ +#include "moe_expert_compute.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#endif + +namespace dflash::common { + +namespace { + +uint64_t hash_u64(uint64_t h, uint64_t v) { + h ^= v + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2); + return h; +} + +const char * nonempty_env(const char * name) { + const char * raw = std::getenv(name); + return raw && *raw ? raw : nullptr; +} + +int parse_nonnegative_env(const char * name, int fallback) { + const char * raw = nonempty_env(name); + if (!raw) return fallback; + errno = 0; + char * end = nullptr; + const long value = std::strtol(raw, &end, 10); + if (errno == ERANGE || end == raw || *end != '\0' || + value < 0 || value > std::numeric_limits::max()) { + return fallback; + } + return (int)value; +} + +std::string make_runtime_key(const MoeExpertComputeRuntimeConfig & cfg) { + std::string key = cfg.target_path; + key += "|n_layer=" + std::to_string(cfg.n_layer); + key += "|n_expert=" + std::to_string(cfg.n_expert); + key += "|n_used=" + std::to_string(cfg.n_expert_used); + key += "|n_embd=" + std::to_string(cfg.n_embd); + key += "|n_ff=" + std::to_string(cfg.n_ff_exp); + if (const char * ipc_bin = nonempty_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_BIN")) { + key += "|ipc_bin="; + key += ipc_bin; + key += "|ipc_gpu=" + std::to_string( + parse_nonnegative_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU", 0)); + key += "|ipc_work="; + if (const char * work_dir = nonempty_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_WORK_DIR")) { + key += work_dir; + } + key += "|ipc_required=" + std::to_string( + parse_nonnegative_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_REQUIRED", 0)); + } else { + key += "|cpu"; + } + return key; +} + +std::string make_multi_target_runtime_key(const MoeExpertComputeRuntimeConfig & cfg, + const ExpertSplitComputeRuntime & split_runtime) { + std::string key = make_runtime_key(cfg); + key += "|multi-target"; + key += "|targets=" + std::to_string(split_runtime.targets.size()); + key += "|split_fp=" + + std::to_string(expert_split_compute_runtime_fingerprint(split_runtime)); + for (const auto & target : split_runtime.targets) { + key += "|"; + key += target.target.name; + key += "="; + key += std::to_string(target.placement.total_hot); + } + return key; +} + +bool validate_executable_file(const char * path, std::string * err) { +#if defined(_WIN32) + (void)path; + (void)err; + return true; +#else + struct stat st {}; + if (::stat(path, &st) != 0) { + if (err) *err = std::string("MoE expert compute IPC binary is not accessible: ") + + path + ": " + std::strerror(errno); + return false; + } + if (!S_ISREG(st.st_mode)) { + if (err) *err = std::string("MoE expert compute IPC binary is not a regular file: ") + path; + return false; + } + if (::access(path, X_OK) != 0) { + if (err) *err = std::string("MoE expert compute IPC binary is not executable: ") + + path + ": " + std::strerror(errno); + return false; + } + return true; +#endif +} + +class MultiTargetMoeExpertCompute final : public MoeExpertCompute { +public: + MultiTargetMoeExpertCompute(MoeMultiTargetExpertRuntime * runtime) + : runtime_(runtime) {} + + bool healthy() const override { + if (!runtime_) return false; + for (const auto & target : runtime_->targets) { + if (!target.compute_active) continue; + if (!target.runtime.compute_ptr() || !target.runtime.compute->healthy()) { + return false; + } + } + return true; + } + + bool prefers_padded_batch() const override { + if (!runtime_) return false; + int active_targets = 0; + bool padded_preferred = false; + for (const auto & target : runtime_->targets) { + if (!target.compute_active) continue; + ++active_targets; + const MoeExpertCompute * compute = target.runtime.compute_ptr(); + if (compute && compute->prefers_padded_batch()) { + padded_preferred = true; + } + } + // A single active target can benefit from padded batches to maximize + // backend-side graph reuse. With multiple active targets, padded dummy + // slots amplify cross-target scatter and can dominate expert_split + // prefill cost, so prefer grouped exact-count batches instead. + return active_targets == 1 && padded_preferred; + } + + bool compute(const MoeExpertLayer & layer, + const float * input, + const int32_t * ids, + const float * weights, + int n_selected, + int n_embd, + int n_ff, + float * output) override { + if (!output || n_selected < 0 || n_embd <= 0 || n_ff <= 0) return false; + if (n_selected > 0 && !input) return false; + if (n_selected == 0) return true; + if (!runtime_ || layer.layer_idx < 0 || + (size_t) layer.layer_idx >= runtime_->layer_routes.size()) { + return false; + } + + const MoeMultiTargetLayerRuntime & layer_rt = + runtime_->layer_routes[(size_t) layer.layer_idx]; + ensure_single_scratch(n_selected, n_embd); + if (!scatter_selected(layer_rt, ids, weights, n_selected)) { + return false; + } + + bool wrote_output = false; + for (size_t ti = 0; ti < runtime_->targets.size(); ++ti) { + const int n_target = target_counts_[(size_t)ti]; + if (n_target <= 0) continue; + auto & target = runtime_->targets[ti]; + const MoeExpertLayer * target_layer = target.runtime.layer_ptr((size_t)layer.layer_idx); + if (!target_layer || !target.runtime.compute_ptr()) { + return false; + } + float * tmp = target_output_scratch_[(size_t)ti].data(); + if (!target.runtime.compute->compute( + *target_layer, + input, + target_ids_[(size_t)ti].data(), + target_weights_[(size_t)ti].data(), + n_target, + n_embd, + n_ff, + tmp)) { + return false; + } + if (!wrote_output) { + std::memcpy(output, tmp, sizeof(float) * (size_t)n_embd); + wrote_output = true; + } else { + for (int i = 0; i < n_embd; ++i) { + output[i] += tmp[i]; + } + } + } + if (!wrote_output) { + std::fill(output, output + n_embd, 0.0f); + } + return true; + } + + bool compute_batch(const MoeExpertLayer & layer, + const float * input, + const int32_t * ids, + const float * weights, + int n_tokens, + int n_selected, + int n_embd, + int n_ff, + float * output) override { + if (!output || n_tokens < 0 || n_selected < 0 || n_embd <= 0 || n_ff <= 0) { + return false; + } + if (n_tokens > 0 && n_selected > 0 && (!input || !ids || !weights)) { + return false; + } + if (n_tokens == 0 || n_selected == 0) { + std::fill(output, output + (size_t)n_tokens * (size_t)n_embd, 0.0f); + return true; + } + if (!runtime_ || layer.layer_idx < 0 || + (size_t) layer.layer_idx >= runtime_->layer_routes.size()) { + return false; + } + + const MoeMultiTargetLayerRuntime & layer_rt = + runtime_->layer_routes[(size_t) layer.layer_idx]; + ensure_batch_scratch(n_tokens, n_selected); + if (!scatter_batch(layer_rt, ids, weights, n_tokens, n_selected)) { + return false; + } + + bool wrote_output = false; + const size_t target_count = runtime_->targets.size(); + for (size_t ti = 0; ti < target_count; ++ti) { + const int max_target_selected = batch_target_max_selected_[(size_t)ti]; + if (max_target_selected <= 0) continue; + + auto & target = runtime_->targets[ti]; + const MoeExpertLayer * target_layer = + target.runtime.layer_ptr((size_t)layer.layer_idx); + MoeExpertCompute * target_compute = target.runtime.compute_ptr(); + if (!target_layer || !target_compute) return false; + + const auto & dense_ids = batch_target_ids_[ti]; + const auto & dense_weights = batch_target_weights_[ti]; + const auto & active_selected_counts = + batch_active_selected_counts_[(size_t)ti]; + for (int n_target : active_selected_counts) { + const auto & token_group = + batch_token_groups_[ti][(size_t)n_target]; + if (token_group.empty()) continue; + if (!dispatch_target_batch( + *target_compute, *target_layer, input, dense_ids, + dense_weights, token_group, n_selected, n_target, + n_embd, n_ff, output, batch_output_written_)) { + return false; + } + wrote_output = true; + } + } + if (!wrote_output) { + std::fill(output, output + (size_t)n_tokens * (size_t)n_embd, 0.0f); + } + + return true; + } + + bool prepare_batch(const MoeExpertLayer & layer, + int n_tokens, + int n_selected, + int n_embd, + int n_ff) override { + if (n_tokens < 0 || n_selected < 0 || n_embd <= 0 || n_ff <= 0) return false; + if (!runtime_ || layer.layer_idx < 0 || + (size_t) layer.layer_idx >= runtime_->layer_routes.size()) { + return false; + } + bool ok = true; + for (auto & target : runtime_->targets) { + if (!target.compute_active) continue; + const MoeExpertLayer * target_layer = target.runtime.layer_ptr((size_t)layer.layer_idx); + if (!target_layer) return false; + const int max_selected = + std::min(n_selected, (int)target_layer->cold_global_by_local.size()); + if (max_selected <= 0) continue; + if (!target.runtime.compute_ptr()) return false; + ok = target.runtime.compute->prepare_batch( + *target_layer, n_tokens, max_selected, n_embd, n_ff) && ok; + } + return ok; + } + + bool prepare_single(const MoeExpertLayer & layer, + int n_selected, + int n_embd, + int n_ff) override { + if (n_selected < 0 || n_embd <= 0 || n_ff <= 0) return false; + if (!runtime_ || layer.layer_idx < 0 || + (size_t) layer.layer_idx >= runtime_->layer_routes.size()) { + return false; + } + bool ok = true; + for (auto & target : runtime_->targets) { + if (!target.compute_active) continue; + const MoeExpertLayer * target_layer = target.runtime.layer_ptr((size_t)layer.layer_idx); + if (!target_layer) return false; + const int max_selected = + std::min(n_selected, (int)target_layer->cold_global_by_local.size()); + if (max_selected <= 0) continue; + if (!target.runtime.compute_ptr()) return false; + ok = target.runtime.compute->prepare_single( + *target_layer, max_selected, n_embd, n_ff) && ok; + } + return ok; + } + +private: + void ensure_single_scratch(int n_selected, int n_embd) { + const size_t target_count = runtime_ ? runtime_->targets.size() : 0; + target_counts_.assign(target_count, 0); + target_ids_.resize(target_count); + target_weights_.resize(target_count); + target_output_scratch_.resize(target_count); + for (size_t ti = 0; ti < target_count; ++ti) { + target_ids_[ti].clear(); + target_weights_[ti].clear(); + if (target_ids_[ti].capacity() < (size_t)n_selected) { + target_ids_[ti].reserve((size_t)n_selected); + } + if (target_weights_[ti].capacity() < (size_t)n_selected) { + target_weights_[ti].reserve((size_t)n_selected); + } + if (target_output_scratch_[ti].size() < (size_t)n_embd) { + target_output_scratch_[ti].resize((size_t)n_embd); + } + } + } + + bool scatter_selected(const MoeMultiTargetLayerRuntime & layer_rt, + const int32_t * ids, + const float * weights, + int n_selected) { + if (!ids || !weights) return false; + const size_t target_count = runtime_ ? runtime_->targets.size() : 0; + for (size_t ti = 0; ti < target_count; ++ti) { + target_counts_[ti] = 0; + target_ids_[ti].clear(); + target_weights_[ti].clear(); + } + + for (int i = 0; i < n_selected; ++i) { + const int32_t union_local = ids[i]; + if (union_local < 0 || + (size_t) union_local >= layer_rt.route_by_union_local.size()) { + return false; + } + const MoeMultiTargetLayerRoute & route = + layer_rt.route_by_union_local[(size_t) union_local]; + if (route.target_slot < 0 || (size_t) route.target_slot >= target_count || + route.target_local < 0) { + return false; + } + target_ids_[(size_t) route.target_slot].push_back(route.target_local); + target_weights_[(size_t) route.target_slot].push_back(weights[i]); + target_counts_[(size_t) route.target_slot]++; + } + return true; + } + + void ensure_batch_scratch(int n_tokens, int n_selected) { + const size_t target_count = runtime_ ? runtime_->targets.size() : 0; + const size_t dense_slots = (size_t)n_tokens * (size_t)n_selected; + batch_target_counts_.resize(target_count); + batch_target_ids_.resize(target_count); + batch_target_weights_.resize(target_count); + batch_token_groups_.resize(target_count); + batch_active_selected_counts_.resize(target_count); + batch_target_max_selected_.assign(target_count, 0); + for (size_t ti = 0; ti < target_count; ++ti) { + auto & token_groups = batch_token_groups_[ti]; + auto & active_selected_counts = + batch_active_selected_counts_[ti]; + for (int n_target : active_selected_counts) { + if (n_target > 0 && (size_t)n_target < token_groups.size()) { + token_groups[(size_t)n_target].clear(); + } + } + active_selected_counts.clear(); + batch_target_counts_[ti].assign((size_t)n_tokens, 0); + batch_target_ids_[ti].resize(dense_slots); + batch_target_weights_[ti].resize(dense_slots); + if (token_groups.size() != (size_t)n_selected + 1) { + token_groups.resize((size_t)n_selected + 1); + } + } + batch_output_written_.assign((size_t)n_tokens, 0); + if (scatter_touched_targets_.capacity() < (size_t)n_selected) { + scatter_touched_targets_.reserve((size_t)n_selected); + } + } + + bool scatter_batch(const MoeMultiTargetLayerRuntime & layer_rt, + const int32_t * ids, + const float * weights, + int n_tokens, + int n_selected) { + const size_t target_count = runtime_ ? runtime_->targets.size() : 0; + for (int t = 0; t < n_tokens; ++t) { + scatter_touched_targets_.clear(); + for (int i = 0; i < n_selected; ++i) { + const size_t src = (size_t)t * (size_t)n_selected + (size_t)i; + const int32_t union_local = ids[src]; + if (union_local < 0 || + (size_t) union_local >= layer_rt.route_by_union_local.size()) { + return false; + } + const MoeMultiTargetLayerRoute & route = + layer_rt.route_by_union_local[(size_t) union_local]; + if (route.target_slot < 0 || + (size_t) route.target_slot >= target_count || + route.target_local < 0) { + return false; + } + auto & counts = batch_target_counts_[(size_t)route.target_slot]; + if (counts[(size_t)t] == 0) { + scatter_touched_targets_.push_back(route.target_slot); + } + const int offset = counts[(size_t)t]++; + if (offset < 0 || offset >= n_selected) return false; + const size_t dst = + (size_t)t * (size_t)n_selected + (size_t)offset; + batch_target_ids_[(size_t)route.target_slot][dst] = + route.target_local; + batch_target_weights_[(size_t)route.target_slot][dst] = + weights[src]; + int & max_selected = + batch_target_max_selected_[(size_t)route.target_slot]; + if (offset + 1 > max_selected) { + max_selected = offset + 1; + } + } + for (int target_slot : scatter_touched_targets_) { + auto & active_selected_counts = + batch_active_selected_counts_[(size_t)target_slot]; + const int count = + batch_target_counts_[(size_t)target_slot][(size_t)t]; + auto pos = std::lower_bound( + active_selected_counts.begin(), + active_selected_counts.end(), + count); + if (pos == active_selected_counts.end() || *pos != count) { + active_selected_counts.insert(pos, count); + } + auto & token_group = + batch_token_groups_[(size_t)target_slot][(size_t)count]; + if (token_group.empty() && + token_group.capacity() < (size_t)n_tokens) { + token_group.reserve((size_t)n_tokens); + } + token_group.push_back(t); + } + } + return true; + } + + bool dispatch_target_batch(MoeExpertCompute & compute, + const MoeExpertLayer & target_layer, + const float * input, + const std::vector & dense_ids, + const std::vector & dense_weights, + const std::vector & token_group, + int n_selected_stride, + int n_target, + int n_embd, + int n_ff, + float * output, + std::vector & output_written) { + const int tc = (int)token_group.size(); + if (tc <= 0) return true; + + group_ids_.resize((size_t)tc * (size_t)n_target); + group_weights_.resize((size_t)tc * (size_t)n_target); + for (int gi = 0; gi < tc; ++gi) { + const int t = token_group[(size_t)gi]; + const size_t src_base = + (size_t)t * (size_t)n_selected_stride; + const size_t dst_base = + (size_t)gi * (size_t)n_target; + for (int i = 0; i < n_target; ++i) { + group_ids_[dst_base + (size_t)i] = + dense_ids[src_base + (size_t)i]; + group_weights_[dst_base + (size_t)i] = + dense_weights[src_base + (size_t)i]; + } + } + + bool contiguous_tokens = true; + for (int gi = 1; gi < tc; ++gi) { + if (token_group[(size_t)gi] != token_group[0] + gi) { + contiguous_tokens = false; + break; + } + } + bool all_unwritten = true; + for (int gi = 0; gi < tc; ++gi) { + const int t = token_group[(size_t)gi]; + if (t < 0 || (size_t)t >= output_written.size()) return false; + if (output_written[(size_t)t]) { + all_unwritten = false; + } + } + + const float * compute_input = input; + if (contiguous_tokens) { + compute_input = input + (size_t)token_group[0] * (size_t)n_embd; + } else { + group_input_.resize((size_t)tc * (size_t)n_embd); + for (int gi = 0; gi < tc; ++gi) { + const int t = token_group[(size_t)gi]; + std::memcpy(group_input_.data() + (size_t)gi * (size_t)n_embd, + input + (size_t)t * (size_t)n_embd, + sizeof(float) * (size_t)n_embd); + } + compute_input = group_input_.data(); + } + + const bool direct_output = contiguous_tokens && all_unwritten; + float * compute_output = direct_output + ? output + (size_t)token_group[0] * (size_t)n_embd + : nullptr; + if (!direct_output) { + group_output_.resize((size_t)tc * (size_t)n_embd); + compute_output = group_output_.data(); + } + if (!compute.compute_batch(target_layer, compute_input, + group_ids_.data(), group_weights_.data(), + tc, n_target, n_embd, n_ff, + compute_output)) { + return false; + } + + if (direct_output) { + for (int gi = 0; gi < tc; ++gi) { + output_written[(size_t)token_group[(size_t)gi]] = 1; + } + return true; + } + + for (int gi = 0; gi < tc; ++gi) { + const int t = token_group[(size_t)gi]; + float * dst = output + (size_t)t * (size_t)n_embd; + const float * src = + group_output_.data() + (size_t)gi * (size_t)n_embd; + if (output_written[(size_t)t]) { + for (int i = 0; i < n_embd; ++i) { + dst[i] += src[i]; + } + } else { + std::memcpy(dst, src, sizeof(float) * (size_t)n_embd); + output_written[(size_t)t] = 1; + } + } + return true; + } + + MoeMultiTargetExpertRuntime * runtime_ = nullptr; + std::vector> target_ids_; + std::vector> target_weights_; + std::vector> target_output_scratch_; + std::vector target_counts_; + std::vector> batch_target_ids_; + std::vector> batch_target_weights_; + std::vector> batch_target_counts_; + std::vector>> batch_token_groups_; + std::vector> batch_active_selected_counts_; + std::vector batch_target_max_selected_; + std::vector batch_output_written_; + std::vector scatter_touched_targets_; + std::vector group_ids_; + std::vector group_weights_; + std::vector group_input_; + std::vector group_output_; +}; + +class RemappedMoeExpertCompute final : public MoeExpertCompute { +public: + RemappedMoeExpertCompute(std::unique_ptr primary, + std::unique_ptr fallback, + std::vector fallback_layers, + std::vector> base_local_by_target_local) + : primary_(std::move(primary)), + fallback_(std::move(fallback)), + fallback_layers_(std::move(fallback_layers)), + base_local_by_target_local_(std::move(base_local_by_target_local)) {} + + bool healthy() const override { + return (primary_ && primary_->healthy()) || + (fallback_ && fallback_->healthy()); + } + + bool prefers_padded_batch() const override { + return primary_ && primary_->prefers_padded_batch(); + } + + bool compute(const MoeExpertLayer & layer, + const float * input, + const int32_t * ids, + const float * weights, + int n_selected, + int n_embd, + int n_ff, + float * output) override { + if (primary_ && + primary_->compute(layer, input, ids, weights, + n_selected, n_embd, n_ff, output)) { + return true; + } + if (!fallback_) return false; + if (!remap_ids(layer.layer_idx, ids, n_selected)) return false; + const MoeExpertLayer * fallback_layer = fallback_layer_ptr(layer.layer_idx); + if (!fallback_layer) return false; + return fallback_->compute(*fallback_layer, input, remapped_ids_.data(), weights, + n_selected, n_embd, n_ff, output); + } + + bool compute_batch(const MoeExpertLayer & layer, + const float * input, + const int32_t * ids, + const float * weights, + int n_tokens, + int n_selected, + int n_embd, + int n_ff, + float * output) override { + if (primary_ && + primary_->compute_batch(layer, input, ids, weights, + n_tokens, n_selected, n_embd, n_ff, output)) { + return true; + } + if (!fallback_) return false; + if (!remap_ids_batch(layer.layer_idx, ids, n_tokens, n_selected)) return false; + const MoeExpertLayer * fallback_layer = fallback_layer_ptr(layer.layer_idx); + if (!fallback_layer) return false; + return fallback_->compute_batch(*fallback_layer, input, remapped_ids_.data(), weights, + n_tokens, n_selected, n_embd, n_ff, output); + } + + bool prepare_batch(const MoeExpertLayer & layer, + int n_tokens, + int n_selected, + int n_embd, + int n_ff) override { + if (primary_ && primary_->prepare_batch(layer, n_tokens, n_selected, n_embd, n_ff)) { + return true; + } + const MoeExpertLayer * fallback_layer = fallback_layer_ptr(layer.layer_idx); + return fallback_ && fallback_layer && + fallback_->prepare_batch(*fallback_layer, n_tokens, n_selected, n_embd, n_ff); + } + + bool prepare_single(const MoeExpertLayer & layer, + int n_selected, + int n_embd, + int n_ff) override { + if (primary_ && primary_->prepare_single(layer, n_selected, n_embd, n_ff)) { + return true; + } + const MoeExpertLayer * fallback_layer = fallback_layer_ptr(layer.layer_idx); + return fallback_ && fallback_layer && + fallback_->prepare_single(*fallback_layer, n_selected, n_embd, n_ff); + } + +private: + const MoeExpertLayer * fallback_layer_ptr(int layer_idx) const { + if (layer_idx < 0 || (size_t)layer_idx >= fallback_layers_.size()) { + return nullptr; + } + return &fallback_layers_[(size_t)layer_idx]; + } + + bool remap_ids(int layer_idx, + const int32_t * ids, + int n_selected) { + if (n_selected < 0 || (n_selected > 0 && !ids) || + layer_idx < 0 || + (size_t)layer_idx >= base_local_by_target_local_.size()) { + return false; + } + const std::vector & lut = base_local_by_target_local_[(size_t)layer_idx]; + remapped_ids_.resize((size_t)n_selected); + for (int i = 0; i < n_selected; ++i) { + const int32_t local = ids[i]; + if (local < 0 || (size_t)local >= lut.size() || lut[(size_t)local] < 0) { + return false; + } + remapped_ids_[(size_t)i] = lut[(size_t)local]; + } + return true; + } + + bool remap_ids_batch(int layer_idx, + const int32_t * ids, + int n_tokens, + int n_selected) { + if (n_tokens < 0 || n_selected < 0 || + ((n_tokens > 0 && n_selected > 0) && !ids) || + layer_idx < 0 || + (size_t)layer_idx >= base_local_by_target_local_.size()) { + return false; + } + const std::vector & lut = base_local_by_target_local_[(size_t)layer_idx]; + remapped_ids_.resize((size_t)n_tokens * (size_t)n_selected); + for (int t = 0; t < n_tokens; ++t) { + for (int i = 0; i < n_selected; ++i) { + const size_t idx = (size_t)t * (size_t)n_selected + (size_t)i; + const int32_t local = ids[idx]; + if (local < 0 || (size_t)local >= lut.size() || lut[(size_t)local] < 0) { + return false; + } + remapped_ids_[idx] = lut[(size_t)local]; + } + } + return true; + } + + std::unique_ptr primary_; + std::unique_ptr fallback_; + std::vector fallback_layers_; + std::vector> base_local_by_target_local_; + std::vector remapped_ids_; +}; + +std::vector build_union_target_layers( + const MoeHybridStorage & hybrid, + const ExpertSplitComputeRuntime & split_runtime, + const std::vector & targets, + std::vector & out_routes) { + std::vector layers(hybrid.layers.size()); + out_routes.assign(hybrid.layers.size(), {}); + for (size_t il = 0; il < hybrid.layers.size(); ++il) { + MoeExpertLayer & layer = layers[(size_t)il]; + const MoeHybridLayerStorage & storage = hybrid.layers[il]; + layer.layer_idx = (int)il; + layer.cold_global_by_local = storage.cold_expert_ids; + out_routes[il].route_by_union_local.resize(layer.cold_global_by_local.size()); + for (size_t union_local = 0; union_local < layer.cold_global_by_local.size(); ++union_local) { + MoeMultiTargetLayerRoute & route = out_routes[il].route_by_union_local[union_local]; + route.target_slot = -1; + route.target_local = -1; + route.union_local = (int)union_local; + } + } + + for (size_t il = 0; il < hybrid.layers.size(); ++il) { + const MoeHybridLayerStorage & storage = hybrid.layers[il]; + MoeMultiTargetLayerRuntime & layer_rt = out_routes[il]; + for (size_t union_local = 0; union_local < storage.cold_expert_ids.size(); ++union_local) { + const int32_t global = storage.cold_expert_ids[union_local]; + if ((int)il >= split_runtime.n_layer || global < 0 || global >= split_runtime.n_expert) { + continue; + } + const int target_slot = split_runtime.target_index((int)il, global); + const int target_local = split_runtime.local_index((int)il, global); + if (target_slot < 0 || target_local < 0 || + (size_t)target_slot >= targets.size()) { + continue; + } + MoeMultiTargetLayerRoute route; + route.target_slot = target_slot; + route.target_local = target_local; + route.union_local = (int)union_local; + layer_rt.route_by_union_local[union_local] = route; + } + } + return layers; +} + +bool parse_target_suffix_device(const std::string & name, + const std::string & prefix, + int & out) { + if (name.size() <= prefix.size() + 1) return false; + if (name.compare(0, prefix.size(), prefix) != 0 || + name[prefix.size()] != ':') { + return false; + } + char * end = nullptr; + long value = std::strtol(name.c_str() + prefix.size() + 1, &end, 10); + if (end == name.c_str() + prefix.size() + 1 || *end != '\0' || value < 0 || + value > std::numeric_limits::max()) { + return false; + } + out = (int)value; + return true; +} + +bool build_target_runtime_from_split( + MoeMultiTargetExpertRuntimeTarget & out, + const MoeExpertComputeRuntimeConfig & cfg, + const ExpertSplitComputeTargetRuntime & split_target, + const MoeHybridStorage & hybrid, + const std::vector & layer_descs, + std::string * err) { + out.target_index = split_target.target_index; + out.target = split_target.target; + out.placement = split_target.placement; + + MoeExpertComputeRuntimeConfig target_cfg = cfg; + const bool is_primary = split_target.target_index == 0; + const bool is_cpu_target = split_target.target.backend == "cpu"; + out.compute_active = + cfg.enabled && !is_primary && split_target.placement.total_hot > 0; + target_cfg.enabled = out.compute_active; + + if (!target_cfg.enabled) { + out.runtime.reset(); + if (is_primary && split_target.placement.total_hot > 0) { + std::fprintf(stderr, + "%s multi-target target_index=%d name=%s backend=%s " + "placement_hot=%d compute=primary-local-skip\n", + cfg.log_prefix ? cfg.log_prefix : "[moe-expert-compute]", + split_target.target_index, + split_target.target.name.c_str(), + split_target.target.backend.c_str(), + split_target.placement.total_hot); + } + return true; + } + + const char * ipc_bin = nonempty_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_BIN"); + const bool remote_required = + parse_nonnegative_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_REQUIRED", 0) != 0; + bool started_remote = false; + std::unique_ptr primary_compute; + if (ipc_bin && !is_cpu_target) { + if (!validate_executable_file(ipc_bin, err)) { + return false; + } + + int target_gpu = 0; + if (const char * gpu_env = nonempty_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU")) { + target_gpu = parse_nonnegative_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU", 0); + (void)gpu_env; + } + if (!split_target.target.backend.empty()) { + if (split_target.target.backend == "cuda") { + (void) parse_target_suffix_device(split_target.target.name, "cuda", target_gpu); + } else if (split_target.target.backend == "hip") { + (void) parse_target_suffix_device(split_target.target.name, "hip", target_gpu); + } + } + + const char * work_dir = nonempty_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_WORK_DIR"); + std::fprintf(stderr, + "%s spawning multi-target expert IPC target_index=%d " + "name=%s backend=%s device=%d placement_hot=%d " + "required=%d reason=secondary-non-cpu\n", + cfg.log_prefix ? cfg.log_prefix : "[moe-expert-compute]", + split_target.target_index, + split_target.target.name.c_str(), + split_target.target.backend.c_str(), + target_gpu, + split_target.placement.total_hot, + remote_required ? 1 : 0); + MoeExpertComputeIpcStartResult remote = + make_moe_expert_compute_ipc_for_placement( + ipc_bin, + cfg.target_path, + target_gpu, + split_target.placement, + cfg.n_embd, + cfg.n_ff_exp, + cfg.n_expert_used, + work_dir ? work_dir : "", + remote_required); + if (remote_required && !remote.started_remote) { + if (err) *err = "required multi-target expert IPC did not start"; + return false; + } + started_remote = remote.started_remote; + if (remote.started_remote) { + primary_compute = std::move(remote.compute); + } + } else if (is_cpu_target && split_target.placement.total_hot > 0) { + std::fprintf(stderr, + "%s multi-target target_index=%d name=%s backend=cpu " + "placement_hot=%d compute=cpu-local\n", + cfg.log_prefix ? cfg.log_prefix : "[moe-expert-compute]", + split_target.target_index, + split_target.target.name.c_str(), + split_target.placement.total_hot); + } + + std::unique_ptr fallback_compute; + std::vector fallback_layers; + std::vector> base_local_by_target_local; + const bool need_cpu_fallback = is_cpu_target || !started_remote || !remote_required; + if (need_cpu_fallback) { + fallback_compute = make_cpu_moe_expert_compute(cfg.n_ff_exp); + fallback_layers = make_moe_expert_layers(hybrid, layer_descs); + base_local_by_target_local.resize((size_t)cfg.n_layer); + for (int il = 0; il < cfg.n_layer; ++il) { + const MoeExpertLayer & base_layer = fallback_layers[(size_t)il]; + const std::vector & globals = + split_target.placement.hot_expert_ids[(size_t)il]; + std::vector remap(globals.size(), -1); + for (size_t local = 0; local < globals.size(); ++local) { + const int32_t global = globals[local]; + if (global < 0 || + (size_t)global >= hybrid.layers[(size_t)il].cold_local_by_global.size()) { + continue; + } + const int32_t base_local = + hybrid.layers[(size_t)il].cold_local_by_global[(size_t)global]; + if (base_local < 0 || + (size_t)base_local >= base_layer.cold_global_by_local.size() || + base_layer.cold_global_by_local[(size_t)base_local] != global) { + continue; + } + remap[local] = base_local; + } + base_local_by_target_local[(size_t)il] = std::move(remap); + } + } + + out.runtime.compute = std::make_unique( + std::move(primary_compute), + std::move(fallback_compute), + std::move(fallback_layers), + std::move(base_local_by_target_local)); + out.runtime.layers.resize((size_t)cfg.n_layer); + for (int il = 0; il < cfg.n_layer; ++il) { + MoeExpertLayer & layer = out.runtime.layers[(size_t)il]; + layer.layer_idx = il; + const auto & globals = split_target.placement.hot_expert_ids[(size_t)il]; + layer.cold_global_by_local.assign(globals.begin(), globals.end()); + } + out.runtime.target_path = cfg.target_path; + out.runtime.runtime_key = + split_target.target.name + (started_remote ? "|ipc" : "|remapped"); + out.runtime.placement_fingerprint = + moe_expert_placement_fingerprint(hybrid, cfg.n_layer, cfg.n_expert, + cfg.n_expert_used) ^ + (uint64_t)(split_target.target_index + 1); + out.runtime.remote_started = started_remote; + return out.runtime.compute_ptr() != nullptr; +} + +} // namespace + +int moe_expert_compute_prepare_batch_limit_from_env() { + const char * raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_PREPARE_BATCH"); + if (!raw || !*raw) return 0; + char * end = nullptr; + long value = std::strtol(raw, &end, 10); + if (end == raw || value <= 0) return 0; + if (value > 4096) return 4096; + return (int)value; +} + +int moe_expert_compute_prepare_selected_limit_from_env() { + const char * raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_PREPARE_SELECTED"); + if (!raw || !*raw) return 0; + char * end = nullptr; + long value = std::strtol(raw, &end, 10); + if (end == raw || value <= 0) return 0; + if (value > 4096) return 4096; + return (int)value; +} + +int moe_expert_compute_batch_limit_from_env() { + const char * raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_BATCH"); + if (!raw || !*raw) return 32; + char * end = nullptr; + long value = std::strtol(raw, &end, 10); + if (end == raw || value <= 0) return 32; + if (value > 4096) return 4096; + return (int)value; +} + +int moe_expert_compute_daemon_batch_limit_from_env() { + const char * raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_DAEMON_BATCH"); + if (!raw || !*raw) return 8; + char * end = nullptr; + long value = std::strtol(raw, &end, 10); + if (end == raw || value <= 0) return 8; + if (value > 4096) return 4096; + return (int)value; +} + +std::unique_ptr make_multi_target_moe_expert_compute( + MoeMultiTargetExpertRuntime * runtime) { + return std::make_unique(runtime); +} + +uint64_t moe_expert_placement_fingerprint(const MoeHybridStorage & hybrid, + int n_layer, + int n_expert, + int n_expert_used) { + uint64_t h = 1469598103934665603ULL; + h = hash_u64(h, (uint64_t)n_layer); + h = hash_u64(h, (uint64_t)n_expert); + h = hash_u64(h, (uint64_t)n_expert_used); + h = hash_u64(h, (uint64_t)hybrid.placement.total_hot); + for (size_t il = 0; il < hybrid.placement.hot_expert_ids.size(); ++il) { + h = hash_u64(h, (uint64_t)il); + for (int32_t expert : hybrid.placement.hot_expert_ids[il]) { + h = hash_u64(h, (uint64_t)(uint32_t)expert); + } + } + return h; +} + +std::vector make_moe_expert_layers( + const MoeHybridStorage & hybrid, + const std::vector & layer_descs) { + std::vector layers(hybrid.layers.size()); + for (size_t il = 0; il < hybrid.layers.size(); ++il) { + const auto & storage = hybrid.layers[il]; + const MoeLayerDesc * desc = + il < layer_descs.size() ? &layer_descs[il] : nullptr; + auto & cl = layers[il]; + cl.layer_idx = (int)il; + cl.cold_global_by_local = storage.cold_expert_ids; + cl.fused_gate_up = (storage.gate_up_cold != nullptr); + if (cl.fused_gate_up) { + cl.gate_up_data = + storage.gate_up_cold ? storage.gate_up_cold->data : nullptr; + cl.gate_up_stride = + storage.gate_up_cold ? storage.gate_up_cold->nb[2] : 0; + cl.gate_up_type = + storage.gate_up_cold ? storage.gate_up_cold->type : GGML_TYPE_Q4_K; + cl.gate_up_scale = desc ? desc->ffn_gate_up_exps_s : 1.0f; + } else { + cl.gate_data = storage.gate_cold ? storage.gate_cold->data : nullptr; + cl.up_data = storage.up_cold ? storage.up_cold->data : nullptr; + cl.gate_stride = storage.gate_cold ? storage.gate_cold->nb[2] : 0; + cl.up_stride = storage.up_cold ? storage.up_cold->nb[2] : 0; + cl.gate_type = + storage.gate_cold ? storage.gate_cold->type : GGML_TYPE_Q4_K; + cl.up_type = + storage.up_cold ? storage.up_cold->type : GGML_TYPE_Q4_K; + cl.gate_scale = desc ? desc->ffn_gate_exps_s : 1.0f; + cl.up_scale = desc ? desc->ffn_up_exps_s : 1.0f; + } + cl.down_data = storage.down_cold ? storage.down_cold->data : nullptr; + cl.down_stride = storage.down_cold ? storage.down_cold->nb[2] : 0; + cl.down_type = + storage.down_cold ? storage.down_cold->type : GGML_TYPE_Q4_K; + cl.down_scale = desc ? desc->ffn_down_exps_s : 1.0f; + } + return layers; +} + +void MoeExpertComputeRuntime::reset() { + compute.reset(); + layers.clear(); + target_path.clear(); + runtime_key.clear(); + placement_fingerprint = 0; + remote_started = false; +} + +void MoeMultiTargetExpertRuntime::reset() { + compute.reset(); + layers.clear(); + targets.clear(); + layer_routes.clear(); + runtime_key.clear(); + placement_fingerprint = 0; + enabled = false; +} + +bool ensure_moe_expert_compute_runtime( + MoeExpertComputeRuntime & runtime, + const MoeExpertComputeRuntimeConfig & cfg, + const MoeHybridStorage & hybrid, + const std::vector & layer_descs, + std::string * err) { + if (!cfg.enabled) { + runtime.reset(); + return true; + } + if (cfg.n_layer <= 0 || cfg.n_expert <= 0 || cfg.n_expert_used <= 0 || + cfg.n_embd <= 0 || cfg.n_ff_exp <= 0) { + if (err) *err = "invalid MoE expert compute runtime config"; + runtime.reset(); + return false; + } + + const uint64_t fingerprint = + moe_expert_placement_fingerprint(hybrid, cfg.n_layer, cfg.n_expert, + cfg.n_expert_used); + const std::string runtime_key = make_runtime_key(cfg); + const bool can_reuse = + runtime.compute && + runtime.compute->healthy() && + runtime.runtime_key == runtime_key && + runtime.placement_fingerprint == fingerprint; + if (!can_reuse) { + runtime.compute.reset(); + runtime.target_path.clear(); + runtime.runtime_key.clear(); + runtime.placement_fingerprint = 0; + runtime.remote_started = false; + } + + if (!runtime.compute) { + bool started_remote = false; + if (const char * ipc_bin = nonempty_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_BIN")) { + if (!validate_executable_file(ipc_bin, err)) { + std::fprintf(stderr, "%s %s\n", cfg.log_prefix ? cfg.log_prefix : "[moe-expert-compute]", + err ? err->c_str() : "invalid remote IPC binary"); + runtime.reset(); + return false; + } + const char * work_dir = nonempty_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_WORK_DIR"); + const int remote_gpu = + parse_nonnegative_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU", 0); + const bool required = + parse_nonnegative_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_REQUIRED", 0) != 0; + MoeExpertComputeIpcStartResult remote = make_moe_expert_compute_ipc( + ipc_bin, cfg.target_path, remote_gpu, hybrid.placement, + cfg.n_embd, cfg.n_ff_exp, cfg.n_expert_used, + work_dir ? work_dir : "", required); + if (required && !remote.started_remote) { + if (err) *err = "remote MoE expert compute IPC is required but did not start"; + std::fprintf(stderr, "%s %s\n", cfg.log_prefix ? cfg.log_prefix : "[moe-expert-compute]", + err ? err->c_str() : "remote IPC did not start"); + runtime.reset(); + return false; + } + started_remote = remote.started_remote; + runtime.compute = std::move(remote.compute); + } + if (!runtime.compute) { + runtime.compute = make_cpu_moe_expert_compute(cfg.n_ff_exp); + } + runtime.remote_started = started_remote; + } + + runtime.layers = make_moe_expert_layers(hybrid, layer_descs); + runtime.target_path = cfg.target_path; + runtime.runtime_key = runtime_key; + runtime.placement_fingerprint = fingerprint; + + if (runtime.remote_started && runtime.compute && runtime.compute->healthy()) { + int prepared = 0; + int failed = 0; + int prepare_tokens = moe_expert_compute_prepare_batch_limit_from_env(); + if (prepare_tokens > 0) { + const bool prepare_full_selected_only = []() { + const char * raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_PREPARE_FULL_SELECTED_ONLY"); + return raw && *raw && std::strcmp(raw, "0") != 0 && + std::strcmp(raw, "false") != 0 && std::strcmp(raw, "off") != 0; + }(); + const int prepare_selected_limit = + moe_expert_compute_prepare_selected_limit_from_env(); + for (const MoeExpertLayer & layer : runtime.layers) { + if (layer.layer_idx < 0 || layer.cold_global_by_local.empty()) continue; + const int max_selected = std::min( + cfg.n_expert_used, (int)layer.cold_global_by_local.size()); + const int first_selected = prepare_full_selected_only + ? max_selected + : 1; + const int last_selected = prepare_full_selected_only + ? max_selected + : (prepare_selected_limit > 0 + ? std::min(max_selected, prepare_selected_limit) + : max_selected); + for (int n_selected = 1; n_selected <= max_selected; ++n_selected) { + if (n_selected < first_selected || n_selected > last_selected) { + continue; + } + const bool prepare_single = + runtime.compute->prepare_single(layer, n_selected, + cfg.n_embd, cfg.n_ff_exp); + const bool prepare_prefill = + prepare_tokens == 1 ? true : runtime.compute->prepare_batch( + layer, prepare_tokens, n_selected, cfg.n_embd, + cfg.n_ff_exp); + if (prepare_single) ++prepared; else ++failed; + if (prepare_tokens != 1) { + if (prepare_prefill) ++prepared; else ++failed; + } + } + } + } + if (prepared > 0 || failed > 0) { + std::fprintf(stderr, + "%s prepared remote MoE graphs=%d failed=%d batch=%d\n", + cfg.log_prefix ? cfg.log_prefix : "[moe-expert-compute]", + prepared, failed, prepare_tokens); + } + } + return true; +} + +bool ensure_multi_target_moe_expert_compute_runtime( + MoeMultiTargetExpertRuntime & runtime, + const MoeExpertComputeRuntimeConfig & cfg, + const ExpertSplitComputeRuntime & split_runtime, + const MoeHybridStorage & hybrid, + const std::vector & layer_descs, + std::string * err) { + if (!cfg.enabled) { + runtime.reset(); + return true; + } + if (!split_runtime.matches(cfg.n_layer, cfg.n_expert, cfg.n_expert_used) || + split_runtime.targets.empty()) { + if (err) *err = "expert split compute runtime not initialized"; + runtime.reset(); + return false; + } + if (split_runtime.targets.size() <= 1) { + runtime.reset(); + return true; + } + + const std::string runtime_key = + make_multi_target_runtime_key(cfg, split_runtime); + const uint64_t fingerprint = + moe_expert_placement_fingerprint(hybrid, cfg.n_layer, cfg.n_expert, + cfg.n_expert_used); + int explicit_non_cpu_targets = 0; + for (const auto & target : split_runtime.targets) { + if (target.target.backend != "cpu") { + ++explicit_non_cpu_targets; + } + } + if (explicit_non_cpu_targets <= 1) { + runtime.reset(); + return true; + } + const bool can_reuse = + runtime.enabled && + runtime.compute && + runtime.compute->healthy() && + runtime.runtime_key == runtime_key && + runtime.placement_fingerprint == fingerprint && + runtime.targets.size() == split_runtime.targets.size(); + if (!can_reuse) { + runtime.reset(); + } + + if (!runtime.compute) { + std::vector targets; + targets.resize(split_runtime.targets.size()); + for (size_t i = 0; i < split_runtime.targets.size(); ++i) { + if (!build_target_runtime_from_split( + targets[i], cfg, split_runtime.targets[i], hybrid, layer_descs, err)) { + runtime.reset(); + return false; + } + } + + std::vector layer_routes; + std::vector union_layers = + build_union_target_layers(hybrid, split_runtime, targets, + layer_routes); + runtime.targets = std::move(targets); + runtime.layer_routes = std::move(layer_routes); + runtime.layers = std::move(union_layers); + runtime.runtime_key = runtime_key; + runtime.placement_fingerprint = fingerprint; + runtime.compute = make_multi_target_moe_expert_compute(&runtime); + runtime.enabled = true; + if (!runtime.compute->healthy()) { + runtime.reset(); + if (err) *err = "multi-target expert compute runtime is unhealthy"; + return false; + } + } + + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/moe_expert_compute.h b/server/src/common/moe_expert_compute.h new file mode 100644 index 000000000..278a3a814 --- /dev/null +++ b/server/src/common/moe_expert_compute.h @@ -0,0 +1,261 @@ +// MoeExpertCompute: direct compute interface for selected MoE expert FFN work. +// +// This is intentionally neutral to hot/cold placement. CPU fallback, remote +// IPC daemons, and future backend-local compute paths should share this shape +// so routing policy can evolve without changing model-specific FFN call sites. +#pragma once + +#include "ggml.h" + +#include +#include +#include +#include +#include +#include + +#include "expert_split_compute_runtime.h" +#include "moe_hybrid_placement.h" +#include "moe_hybrid_storage.h" + +namespace dflash::common { + +// Per-layer selected expert weight metadata: raw pointers into the placement +// storage for whichever experts this compute implementation owns. +struct MoeExpertLayer { + int layer_idx = -1; + std::vector cold_global_by_local; + + const void * gate_up_data = nullptr; // fused [n_expert, n_ff*2, n_embd] + const void * gate_data = nullptr; // separate gate [n_expert, n_ff, n_embd] + const void * up_data = nullptr; // separate up [n_expert, n_ff, n_embd] + const void * down_data = nullptr; // [n_expert, n_embd, n_ff] + + size_t gate_up_stride = 0; + size_t gate_stride = 0; + size_t up_stride = 0; + size_t down_stride = 0; + + ggml_type gate_up_type = GGML_TYPE_Q4_K; + ggml_type gate_type = GGML_TYPE_Q4_K; + ggml_type up_type = GGML_TYPE_Q4_K; + ggml_type down_type = GGML_TYPE_Q4_K; + bool fused_gate_up = false; + + float gate_up_scale = 1.0f; + float gate_scale = 1.0f; + float up_scale = 1.0f; + float down_scale = 1.0f; +}; + +struct MoeExpertCompute { + virtual ~MoeExpertCompute() = default; + virtual bool healthy() const { return true; } + virtual bool prefers_padded_batch() const { return false; } + + // Compute selected expert FFN contributions and accumulate into output. + // input: [n_embd] F32, post-norm hidden state + // ids: [n_selected] I32, local expert indices for this placement + // weights: [n_selected] F32, routing weights + // output: [n_embd] F32, accumulated weighted expert output + virtual bool compute( + const MoeExpertLayer & layer, + const float * input, + const int32_t * ids, + const float * weights, + int n_selected, + int n_embd, + int n_ff, + float * output) = 0; + + // Optional remote-graph preparation hook. Implementations that keep stable + // backend-side graph state can prebuild common batch shapes before the + // first request payload arrives. Local/CPU implementations can no-op. + virtual bool prepare_batch( + const MoeExpertLayer & layer, + int n_tokens, + int n_selected, + int n_embd, + int n_ff) { + (void)layer; + (void)n_ff; + return n_tokens >= 0 && n_selected >= 0 && n_embd > 0; + } + + // Optional decode-path preparation hook. This is the single-token companion + // to prepare_batch() for remote implementations that keep stable decode + // graphs alive on the backend. + virtual bool prepare_single( + const MoeExpertLayer & layer, + int n_selected, + int n_embd, + int n_ff) { + (void)layer; + (void)n_ff; + return n_selected >= 0 && n_embd > 0; + } + + // Coarse-grained prefill hook. Remote implementations override this to + // amortize command/payload overhead; CPU fallback keeps the old per-token + // behavior through the single-token compute path. + virtual bool compute_batch( + const MoeExpertLayer & layer, + const float * input, + const int32_t * ids, + const float * weights, + int n_tokens, + int n_selected, + int n_embd, + int n_ff, + float * output) { + if (n_tokens < 0 || n_selected < 0 || n_embd <= 0 || !output) return false; + if (n_tokens == 0 || n_selected == 0) { + std::fill(output, output + (size_t)n_tokens * (size_t)n_embd, 0.0f); + return true; + } + if (!input || !ids || !weights) return false; + for (int t = 0; t < n_tokens; ++t) { + if (!compute(layer, + input + (size_t)t * (size_t)n_embd, + ids + (size_t)t * (size_t)n_selected, + weights + (size_t)t * (size_t)n_selected, + n_selected, + n_embd, + n_ff, + output + (size_t)t * (size_t)n_embd)) { + return false; + } + } + return true; + } +}; + +std::unique_ptr make_cpu_moe_expert_compute(int n_ff_max); +int moe_expert_compute_batch_limit_from_env(); +int moe_expert_compute_prepare_batch_limit_from_env(); +int moe_expert_compute_prepare_selected_limit_from_env(); +int moe_expert_compute_daemon_batch_limit_from_env(); + +struct MoeExpertComputeIpcStartResult { + std::unique_ptr compute; + bool started_remote = false; +}; + +struct MoeExpertComputeRuntime { + std::unique_ptr compute; + std::vector layers; + std::string target_path; + std::string runtime_key; + uint64_t placement_fingerprint = 0; + bool remote_started = false; + + void reset(); + MoeExpertCompute * compute_ptr() const { return compute.get(); } + const MoeExpertLayer * layer_ptr(size_t il) const { + return il < layers.size() ? &layers[il] : nullptr; + } +}; + +struct MoeMultiTargetExpertRuntimeTarget { + int target_index = -1; + ExpertSplitTarget target; + MoeHybridPlacement placement; + MoeExpertComputeRuntime runtime; + bool compute_active = false; +}; + +struct MoeMultiTargetLayerRoute { + int target_slot = -1; + int target_local = -1; + int union_local = -1; +}; + +struct MoeMultiTargetLayerRuntime { + std::vector route_by_union_local; +}; + +struct MoeMultiTargetExpertRuntime { + std::unique_ptr compute; + std::vector layers; + std::vector targets; + std::vector layer_routes; + std::string runtime_key; + uint64_t placement_fingerprint = 0; + bool enabled = false; + + void reset(); + MoeExpertCompute * compute_ptr() const { return compute.get(); } + const MoeExpertLayer * layer_ptr(size_t il) const { + return il < layers.size() ? &layers[il] : nullptr; + } +}; + +struct MoeExpertComputeRuntimeConfig { + std::string target_path; + int n_layer = 0; + int n_expert = 0; + int n_expert_used = 0; + int n_embd = 0; + int n_ff_exp = 0; + bool enabled = true; + const char * log_prefix = "[moe-expert-compute]"; +}; + +std::unique_ptr make_multi_target_moe_expert_compute( + MoeMultiTargetExpertRuntime * runtime); + +uint64_t moe_expert_placement_fingerprint(const MoeHybridStorage & hybrid, + int n_layer, + int n_expert, + int n_expert_used); + +std::vector make_moe_expert_layers( + const MoeHybridStorage & hybrid, + const std::vector & layer_descs); + +bool ensure_moe_expert_compute_runtime( + MoeExpertComputeRuntime & runtime, + const MoeExpertComputeRuntimeConfig & cfg, + const MoeHybridStorage & hybrid, + const std::vector & layer_descs, + std::string * err = nullptr); + +bool ensure_multi_target_moe_expert_compute_runtime( + MoeMultiTargetExpertRuntime & runtime, + const MoeExpertComputeRuntimeConfig & cfg, + const ExpertSplitComputeRuntime & split_runtime, + const MoeHybridStorage & hybrid, + const std::vector & layer_descs, + std::string * err = nullptr); + +MoeExpertComputeIpcStartResult make_moe_expert_compute_ipc( + const std::string & bin, + const std::string & target_path, + int target_gpu, + const MoeHybridPlacement & main_placement, + int n_embd, + int n_ff_exp, + int n_expert_used, + const std::string & work_dir, + bool required); + +MoeExpertComputeIpcStartResult make_moe_expert_compute_ipc_for_placement( + const std::string & bin, + const std::string & target_path, + int target_gpu, + const MoeHybridPlacement & remote_placement, + int n_embd, + int n_ff_exp, + int n_expert_used, + const std::string & work_dir, + bool required); + +int run_moe_expert_compute_ipc_daemon(const char * target_path, + const char * placement_path, + int target_gpu, + int stream_fd, + int payload_fd = -1, + int shared_payload_fd = -1, + size_t shared_payload_bytes = 0); + +} // namespace dflash::common diff --git a/server/src/common/moe_expert_compute_cpu.cpp b/server/src/common/moe_expert_compute_cpu.cpp new file mode 100644 index 000000000..edbf52960 --- /dev/null +++ b/server/src/common/moe_expert_compute_cpu.cpp @@ -0,0 +1,195 @@ +// CpuMoeExpertCompute: fused selected-expert FFN using ggml vec_dot primitives. +// Bypasses ggml graph dispatch overhead. Uses OpenMP to saturate memory bandwidth. + +#include "moe_expert_compute.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include + +#ifdef _OPENMP +#include +#endif + +namespace dflash::common { + +class CpuMoeExpertCompute : public MoeExpertCompute { + int n_ff_max_; + int n_threads_; + + // Per-thread scratch buffers for parallel down matmul + struct ThreadBuf { + std::vector scratch; // [n_ff * 2] gate_up result + SwiGLU + std::vector mid_conv; // down input converted to vec_dot_type + }; + std::vector thread_bufs_; + std::vector inp_conv_; // input converted (shared, read-only during matmul) + +public: + explicit CpuMoeExpertCompute(int n_ff_max, int n_threads = 0) : n_ff_max_(n_ff_max) { +#ifdef _OPENMP + if (n_threads <= 0) { + const char * env = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_THREADS"); + if (!env || !*env) env = std::getenv("DFLASH_COLD_THREADS"); + n_threads = env ? std::atoi(env) : 0; + } + n_threads_ = n_threads > 0 ? n_threads : std::min(omp_get_max_threads(), 8); +#else + n_threads_ = 1; +#endif + fprintf(stderr, "[moe-expert-compute] cpu threads=%d\n", n_threads_); + thread_bufs_.resize(n_threads_); + for (auto & tb : thread_bufs_) { + tb.scratch.resize((size_t)n_ff_max * 2); + } + } + + bool compute( + const MoeExpertLayer & layer, + const float * input, + const int32_t * ids, + const float * weights, + int n_selected, + int n_embd, + int n_ff, + float * output) override { + + if (n_selected < 0 || n_embd <= 0 || n_ff <= 0 || !input || !ids || + !weights || !output) { + return false; + } + std::memset(output, 0, sizeof(float) * (size_t)n_embd); + if (n_selected == 0) return true; + + // Gate/up phase type traits + const ggml_type gu_type = layer.fused_gate_up ? layer.gate_up_type : layer.gate_type; + const auto * gu_cpu_traits = ggml_get_type_traits_cpu(gu_type); + const auto gu_vec_dot = gu_cpu_traits->vec_dot; + const auto gu_vec_dot_type = gu_cpu_traits->vec_dot_type; + const auto gu_from_float = ggml_get_type_traits_cpu(gu_vec_dot_type)->from_float; + + // Down phase type traits (may differ from gate/up) + const auto * dn_cpu_traits = ggml_get_type_traits_cpu(layer.down_type); + const auto dn_vec_dot = dn_cpu_traits->vec_dot; + const auto dn_vec_dot_type = dn_cpu_traits->vec_dot_type; + const auto dn_from_float = ggml_get_type_traits_cpu(dn_vec_dot_type)->from_float; + + const size_t inp_row_size = ggml_row_size(gu_vec_dot_type, n_embd); + const size_t mid_row_size = ggml_row_size(dn_vec_dot_type, n_ff); + const size_t gu_weight_row = ggml_row_size(gu_type, n_embd); + const size_t dn_weight_row = ggml_row_size(layer.down_type, n_ff); + + // For separate gate/up — up may have a different type than gate + size_t up_weight_row = gu_weight_row; + const ggml_type up_type_actual = layer.fused_gate_up ? gu_type : layer.up_type; + (void)up_type_actual; + ggml_vec_dot_t up_vec_dot = gu_vec_dot; + ggml_type up_vdt = gu_vec_dot_type; + if (!layer.fused_gate_up && layer.up_type != layer.gate_type) { + const auto * up_cpu_traits = ggml_get_type_traits_cpu(layer.up_type); + up_vec_dot = up_cpu_traits->vec_dot; + up_vdt = up_cpu_traits->vec_dot_type; + up_weight_row = ggml_row_size(layer.up_type, n_embd); + } + + // Ensure input conversion buffer is large enough + if (inp_conv_.size() < inp_row_size) inp_conv_.resize(inp_row_size); + // Ensure per-thread mid_conv buffers + for (auto & tb : thread_bufs_) { + if (tb.mid_conv.size() < mid_row_size) tb.mid_conv.resize(mid_row_size); + } + + // Convert input for up if different type + std::vector inp_conv_up; + if (!layer.fused_gate_up && up_vdt != gu_vec_dot_type) { + size_t up_inp_row_size = ggml_row_size(up_vdt, n_embd); + inp_conv_up.resize(up_inp_row_size); + auto up_from_float = ggml_get_type_traits_cpu(up_vdt)->from_float; + up_from_float(input, inp_conv_up.data(), n_embd); + } + + // Convert input to gate's vec_dot format once + gu_from_float(input, inp_conv_.data(), n_embd); + + for (int e = 0; e < n_selected; ++e) { + const int32_t eid = ids[e]; + const float w = weights[e]; + if (w == 0.0f) continue; + + // Use thread 0's scratch for gate_up (serial phase) + float * scratch = thread_bufs_[0].scratch.data(); + + // ── Phase 1: gate_up matmul → scratch[0..n_ff*2) ── + // Parallel over rows (each row is independent, reading shared inp_conv_) + if (layer.fused_gate_up) { + const char * expert = (const char *)layer.gate_up_data + (size_t)eid * layer.gate_up_stride; + const int n_rows = n_ff * 2; +#ifdef _OPENMP + #pragma omp parallel for num_threads(n_threads_) schedule(static) +#endif + for (int row = 0; row < n_rows; ++row) { + const void * row_ptr = expert + (size_t)row * gu_weight_row; + gu_vec_dot(n_embd, &scratch[row], 0, row_ptr, 0, inp_conv_.data(), 0, 1); + } + if (layer.gate_up_scale != 1.0f) { + for (int i = 0; i < n_rows; ++i) scratch[i] *= layer.gate_up_scale; + } + } else { + const char * gate_expert = (const char *)layer.gate_data + (size_t)eid * layer.gate_stride; + const char * up_expert = (const char *)layer.up_data + (size_t)eid * layer.up_stride; + const uint8_t * up_inp = (!inp_conv_up.empty()) ? inp_conv_up.data() : inp_conv_.data(); +#ifdef _OPENMP + #pragma omp parallel for num_threads(n_threads_) schedule(static) +#endif + for (int row = 0; row < n_ff; ++row) { + const void * gp = gate_expert + (size_t)row * gu_weight_row; + gu_vec_dot(n_embd, &scratch[row], 0, gp, 0, inp_conv_.data(), 0, 1); + const void * up = up_expert + (size_t)row * up_weight_row; + up_vec_dot(n_embd, &scratch[n_ff + row], 0, up, 0, up_inp, 0, 1); + } + if (layer.gate_scale != 1.0f) { + for (int i = 0; i < n_ff; ++i) scratch[i] *= layer.gate_scale; + } + if (layer.up_scale != 1.0f) { + for (int i = 0; i < n_ff; ++i) scratch[n_ff + i] *= layer.up_scale; + } + } + + // ── Phase 2: SwiGLU activation ── + for (int i = 0; i < n_ff; ++i) { + const float gate = scratch[i]; + const float up = scratch[n_ff + i]; + scratch[i] = (gate / (1.0f + expf(-gate))) * up; + } + + // ── Phase 3: down matmul → output (weighted accumulate) ── + // Convert SwiGLU result to down's vec_dot format (serial, small) + dn_from_float(scratch, thread_bufs_[0].mid_conv.data(), n_ff); + const uint8_t * mid_conv_data = thread_bufs_[0].mid_conv.data(); + + const char * down_expert = (const char *)layer.down_data + (size_t)eid * layer.down_stride; + const float scale = w * layer.down_scale; + + // Parallel down matmul — each thread accumulates its own output rows +#ifdef _OPENMP + #pragma omp parallel for num_threads(n_threads_) schedule(static) +#endif + for (int row = 0; row < n_embd; ++row) { + float val; + const void * row_ptr = down_expert + (size_t)row * dn_weight_row; + dn_vec_dot(n_ff, &val, 0, row_ptr, 0, mid_conv_data, 0, 1); + output[row] += scale * val; + } + } + return true; + } +}; + +std::unique_ptr make_cpu_moe_expert_compute(int n_ff_max) { + return std::make_unique(n_ff_max); +} + +} // namespace dflash::common diff --git a/server/src/common/moe_expert_compute_ipc.cpp b/server/src/common/moe_expert_compute_ipc.cpp new file mode 100644 index 000000000..33bb278a2 --- /dev/null +++ b/server/src/common/moe_expert_compute_ipc.cpp @@ -0,0 +1,2075 @@ +// Mixed-backend MoE expert compute IPC. +// +// The parent backend computes router + hot experts. This client sends the +// selected non-local expert work to a backend-local daemon, which computes only +// that routed expert partial and returns one F32 hidden vector. + +#include "moe_expert_compute.h" + +#include "backend_ipc.h" +#include "io_utils.h" +#include "moe_hybrid_ffn_eval.h" +#include "moe_hybrid_placement.h" +#include "moe_hybrid_storage.h" +#include "gguf_tensor_data.h" +#include "ggml_graph_precision.h" +#include "internal.h" +#include "laguna_internal.h" +#include "moe_hybrid_types_impl.h" + +#include "ggml-backend.h" +#include "ggml-cuda.h" +#include "gguf.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +# include +# include +# include +# include +#endif + +namespace dflash::common { + +namespace { + +using MoeExpertClock = std::chrono::steady_clock; + +uint64_t moe_expert_compute_elapsed_us(MoeExpertClock::time_point start, + MoeExpertClock::time_point end) { + return (uint64_t)std::chrono::duration_cast( + end - start).count(); +} + +bool moe_expert_compute_profile_enabled() { + const char * raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_IPC_PROFILE"); + return raw && *raw && std::strcmp(raw, "0") != 0 && + std::strcmp(raw, "false") != 0 && std::strcmp(raw, "off") != 0; +} + +#if !defined(_WIN32) +bool moe_expert_compute_drain_exact_fd(int fd, size_t bytes) { + std::array tmp{}; + while (bytes > 0) { + const size_t chunk = std::min(bytes, tmp.size()); + if (!read_exact_fd(fd, tmp.data(), chunk)) return false; + bytes -= chunk; + } + return true; +} +#endif + +bool moe_expert_compute_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; +} + +struct MoeExpertIpcClientStats { + uint64_t calls = 0; + uint64_t payload_bytes = 0; + uint64_t response_bytes = 0; + uint64_t pack_us = 0; + uint64_t send_us = 0; + uint64_t wait_us = 0; + uint64_t recv_us = 0; + + void maybe_print(bool force) const { + if (!force && calls > 0 && calls % 128 != 0) return; + if (calls == 0) return; + std::fprintf(stderr, + "[moe-expert-compute-ipc] profile calls=%" PRIu64 + " payload_mib=%.3f response_mib=%.3f" + " pack_ms=%.3f send_ms=%.3f wait_ms=%.3f recv_ms=%.3f\n", + calls, payload_bytes / 1048576.0, + response_bytes / 1048576.0, + pack_us / 1000.0, send_us / 1000.0, + wait_us / 1000.0, recv_us / 1000.0); + } +}; + +struct MoeExpertDaemonStats { + uint64_t calls = 0; + uint64_t payload_bytes = 0; + uint64_t response_bytes = 0; + uint64_t read_us = 0; + uint64_t unpack_us = 0; + uint64_t graph_builds = 0; + uint64_t graph_reuses = 0; + uint64_t graph_build_us = 0; + uint64_t tensor_set_us = 0; + uint64_t compute_us = 0; + uint64_t tensor_get_us = 0; + uint64_t write_us = 0; + + void maybe_print(bool force, const char * scope = "") const { + if (!force && calls > 0 && calls % 128 != 0) return; + if (calls == 0) return; + std::fprintf(stderr, + "[moe-expert-compute-daemon] profile%s calls=%" PRIu64 + " payload_mib=%.3f response_mib=%.3f builds=%" PRIu64 + " reuses=%" PRIu64 + " read_ms=%.3f unpack_ms=%.3f build_ms=%.3f set_ms=%.3f" + " compute_ms=%.3f get_ms=%.3f write_ms=%.3f\n", + scope, calls, payload_bytes / 1048576.0, + response_bytes / 1048576.0, graph_builds, graph_reuses, + read_us / 1000.0, unpack_us / 1000.0, + graph_build_us / 1000.0, tensor_set_us / 1000.0, + compute_us / 1000.0, tensor_get_us / 1000.0, + write_us / 1000.0); + } +}; + +inline ggml_tensor * moe_expert_apply_scale2(ggml_context * ctx, + ggml_tensor * mm_result, + float scale) { + return scale == 1.0f ? mm_result : ggml_scale(ctx, mm_result, scale); +} + +bool moe_expert_ipc_is_supported_input_type(ggml_type type); + +struct CachedBatchedMoeExpertGraph { + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t alloc = nullptr; + ggml_tensor * inp = nullptr; + ggml_tensor * ids = nullptr; + ggml_tensor * weights = nullptr; + ggml_tensor * output = nullptr; + int n_tokens = 0; + int n_selected = 0; + ggml_type input_type = GGML_TYPE_F32; + + bool valid() const { return ctx && gf && alloc && output; } + + void free() { + if (alloc) { ggml_gallocr_free(alloc); alloc = nullptr; } + if (ctx) { ggml_free(ctx); ctx = nullptr; } + gf = nullptr; + inp = nullptr; + ids = nullptr; + weights = nullptr; + output = nullptr; + n_tokens = 0; + n_selected = 0; + input_type = GGML_TYPE_F32; + } +}; + +struct RemoteMoeLayerRuntime { + float gate_scale = 1.0f; + float up_scale = 1.0f; + float down_scale = 1.0f; + float gate_up_scale = 1.0f; +}; + +struct RemoteMoeRuntime { + std::string arch; + int n_layer = 0; + int n_embd = 0; + int n_expert = 0; + int n_expert_used = 0; + int n_ff_exp = 0; + std::vector layers; + MoeHybridStorage hybrid; +}; + +bool build_cached_batched_cold_graph( + CachedBatchedMoeExpertGraph & out, + ggml_backend_t backend, + ggml_tensor * gate_tensor, + ggml_tensor * up_tensor, + ggml_tensor * down_tensor, + ggml_tensor * gate_up_tensor, + float gate_scale, + float up_scale, + float down_scale, + float gate_up_scale, + int n_embd, + int n_ff_exp, + int n_selected, + int n_tokens, + ggml_type input_type = GGML_TYPE_F32) { + + out.free(); + if (n_embd <= 0 || n_ff_exp <= 0 || n_selected <= 0 || n_tokens <= 0 || + !down_tensor || (!gate_up_tensor && (!gate_tensor || !up_tensor)) || + !moe_expert_ipc_is_supported_input_type(input_type)) { + return false; + } + + ggml_init_params ip{}; + ip.mem_size = 128 * 1024 * 1024; + ip.mem_buffer = nullptr; + ip.no_alloc = true; + out.ctx = ggml_init(ip); + if (!out.ctx) return false; + + out.inp = ggml_new_tensor_2d(out.ctx, input_type, n_embd, n_tokens); + ggml_set_input(out.inp); + out.ids = ggml_new_tensor_2d(out.ctx, GGML_TYPE_I32, n_selected, n_tokens); + ggml_set_input(out.ids); + out.weights = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, n_selected, n_tokens); + ggml_set_input(out.weights); + + ggml_tensor * cur_f32 = graph_tensor_f32(out.ctx, out.inp); + ggml_tensor * cur_3d = ggml_reshape_3d(out.ctx, cur_f32, n_embd, 1, n_tokens); + ggml_tensor * gu = nullptr; + if (gate_up_tensor) { + ggml_tensor * gate_up_e = moe_expert_apply_scale2(out.ctx, + ggml_mul_mat_id(out.ctx, gate_up_tensor, cur_3d, out.ids), + gate_up_scale); + ggml_tensor * gate_e = ggml_view_3d(out.ctx, gate_up_e, + n_ff_exp, gate_up_e->ne[1], gate_up_e->ne[2], + gate_up_e->nb[1], gate_up_e->nb[2], 0); + ggml_tensor * up_e = ggml_view_3d(out.ctx, gate_up_e, + n_ff_exp, gate_up_e->ne[1], gate_up_e->ne[2], + gate_up_e->nb[1], gate_up_e->nb[2], + (size_t)n_ff_exp * ggml_element_size(gate_up_e)); + gate_e = ggml_cont(out.ctx, gate_e); + up_e = ggml_cont(out.ctx, up_e); + gu = ggml_swiglu_split(out.ctx, gate_e, up_e); + } else { + ggml_tensor * gate_e = moe_expert_apply_scale2(out.ctx, + ggml_mul_mat_id(out.ctx, gate_tensor, cur_3d, out.ids), + gate_scale); + ggml_tensor * up_e = moe_expert_apply_scale2(out.ctx, + ggml_mul_mat_id(out.ctx, up_tensor, cur_3d, out.ids), + up_scale); + gu = ggml_swiglu_split(out.ctx, gate_e, up_e); + } + + ggml_tensor * experts = moe_expert_apply_scale2(out.ctx, + ggml_mul_mat_id(out.ctx, down_tensor, gu, out.ids), down_scale); + ggml_tensor * w_view = ggml_reshape_3d(out.ctx, out.weights, + 1, n_selected, n_tokens); + experts = ggml_mul(out.ctx, experts, w_view); + + ggml_tensor * sum_shape = + ggml_new_tensor_3d(out.ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); + ggml_tensor * moe_sum = ggml_repeat_back(out.ctx, experts, sum_shape); + out.output = ggml_reshape_2d(out.ctx, moe_sum, n_embd, n_tokens); + out.gf = ggml_new_graph_custom(out.ctx, 4096, false); + ggml_set_output(out.output); + ggml_build_forward_expand(out.gf, out.output); + out.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(out.alloc, out.gf)) { + out.free(); + return false; + } + out.n_tokens = n_tokens; + out.n_selected = n_selected; + out.input_type = input_type; + return true; +} + +bool emit_status(int stream_fd, int32_t status) { + if (stream_fd < 0) return false; +#if defined(_WIN32) + stream_emit_fd(stream_fd, status); + return true; +#else + return write_exact_fd(stream_fd, &status, sizeof(status)); +#endif +} + +uint64_t moe_expert_compute_batch_graph_key(int n_tokens, int n_selected) { + return ((uint64_t)(uint32_t)n_tokens << 32) | (uint32_t)n_selected; +} + +uint64_t moe_expert_compute_batch_graph_key(int n_tokens, + int n_selected, + ggml_type input_type) { + uint64_t key = ((uint64_t)(uint32_t)n_tokens << 32) | + ((uint64_t)(uint32_t)n_selected << 8) | + (uint64_t)(uint8_t)input_type; + return key; +} + +MoeHybridPlacement invert_moe_hybrid_placement(const MoeHybridPlacement & main) { + MoeHybridPlacement remote; + remote.n_layer = main.n_layer; + remote.n_expert = main.n_expert; + remote.n_expert_used = main.n_expert_used; + remote.hot_counts.assign((size_t)main.n_layer, 0); + remote.hot_expert_ids.assign((size_t)main.n_layer, {}); + remote.total_hot = 0; + + for (int il = 0; il < main.n_layer; ++il) { + std::vector is_main_hot((size_t)main.n_expert, 0); + if ((size_t)il < main.hot_expert_ids.size()) { + for (int32_t expert : main.hot_expert_ids[(size_t)il]) { + if (expert >= 0 && expert < main.n_expert) { + is_main_hot[(size_t)expert] = 1; + } + } + } + auto & remote_hot = remote.hot_expert_ids[(size_t)il]; + for (int expert = 0; expert < main.n_expert; ++expert) { + if (!is_main_hot[(size_t)expert]) { + remote_hot.push_back((int32_t)expert); + } + } + remote.hot_counts[(size_t)il] = (int)remote_hot.size(); + remote.total_hot += (int)remote_hot.size(); + } + return remote; +} + +bool write_temp_remote_placement(const MoeHybridPlacement & main, + std::string & path_out, + std::string * err) { +#if defined(_WIN32) + (void)main; (void)path_out; + if (err) *err = "remote MoE expert compute IPC is POSIX-only"; + return false; +#else + char templ[] = "/tmp/lucebox_moe_cold_placement_XXXXXX"; + int fd = ::mkstemp(templ); + if (fd < 0) { + if (err) *err = std::string("mkstemp failed: ") + std::strerror(errno); + return false; + } + ::close(fd); + path_out = templ; + const MoeHybridPlacement remote = invert_moe_hybrid_placement(main); + if (!remote.save_json(path_out, "moe_remote_expert_compute", err)) { + ::unlink(path_out.c_str()); + path_out.clear(); + return false; + } + return true; +#endif +} + +bool write_temp_placement_file(const MoeHybridPlacement & placement, + const char * arch_name, + std::string & path_out, + std::string * err) { +#if defined(_WIN32) + (void)placement; + (void)arch_name; + (void)path_out; + if (err) *err = "remote MoE expert compute IPC is POSIX-only"; + return false; +#else + char templ[] = "/tmp/lucebox_moe_target_placement_XXXXXX"; + int fd = ::mkstemp(templ); + if (fd < 0) { + if (err) *err = std::string("mkstemp failed: ") + std::strerror(errno); + return false; + } + ::close(fd); + path_out = templ; + if (!placement.save_json(path_out, arch_name ? arch_name : "moe_expert_target", err)) { + ::unlink(path_out.c_str()); + path_out.clear(); + return false; + } + return true; +#endif +} + +bool write_named_placement_file(const MoeHybridPlacement & placement, + const std::string & path, + const char * arch_name, + std::string * err) { + return placement.save_json(path, arch_name ? arch_name : "moe_expert_target", err); +} + +BackendIpcPayloadTransport moe_expert_transport_from_env() { + const char * raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_IPC_TRANSPORT"); + if (!raw || !*raw) { + return BackendIpcPayloadTransport::Auto; + } + BackendIpcPayloadTransport transport = BackendIpcPayloadTransport::Stream; + if (!parse_backend_ipc_payload_transport(raw, transport)) { + return BackendIpcPayloadTransport::Auto; + } + return transport; +} + +int moe_expert_ipc_shared_batch_capacity_from_env() { + const char * raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_IPC_BATCH_CAPACITY"); + if (!raw || !*raw) return 1024; + char * end = nullptr; + long value = std::strtol(raw, &end, 10); + if (end == raw || value <= 0) return 1024; + if (value > 4096) return 4096; + return (int)value; +} + +ggml_type moe_expert_ipc_input_type_from_env() { + const char * raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_IPC_DTYPE"); + if (!raw || !*raw || std::strcmp(raw, "f32") == 0 || std::strcmp(raw, "F32") == 0) { + return GGML_TYPE_F32; + } + if (std::strcmp(raw, "f16") == 0 || std::strcmp(raw, "F16") == 0) { + return GGML_TYPE_F16; + } + if (std::strcmp(raw, "bf16") == 0 || std::strcmp(raw, "BF16") == 0) { + return GGML_TYPE_BF16; + } + std::fprintf(stderr, + "[moe-expert-compute-ipc] ignoring unsupported " + "DFLASH_MOE_EXPERT_COMPUTE_IPC_DTYPE=%s\n", + raw); + return GGML_TYPE_F32; +} + +bool moe_expert_ipc_is_supported_input_type(ggml_type type) { + return type == GGML_TYPE_F32 || type == GGML_TYPE_F16 || + type == GGML_TYPE_BF16; +} + +size_t moe_expert_ipc_input_row_size(ggml_type type, size_t n_embd) { + if (!moe_expert_ipc_is_supported_input_type(type) || n_embd == 0) { + return 0; + } + return ggml_row_size(type, (int64_t)n_embd); +} + +bool moe_expert_convert_input_to_ipc_type(ggml_type type, + const float * src, + int n_tokens, + int n_embd, + std::vector & scratch, + const void *& data, + size_t & bytes) { + data = src; + bytes = 0; + if (!src || n_tokens <= 0 || n_embd <= 0 || + !moe_expert_ipc_is_supported_input_type(type)) { + return false; + } + const size_t elems = (size_t)n_tokens * (size_t)n_embd; + if (type == GGML_TYPE_F32) { + bytes = elems * sizeof(float); + return true; + } + bytes = ggml_row_size(type, (int64_t)n_embd) * (size_t)n_tokens; + scratch.resize(bytes); + if (type == GGML_TYPE_F16) { + ggml_fp32_to_fp16_row(src, reinterpret_cast(scratch.data()), + (int64_t)elems); + } else { + ggml_fp32_to_bf16_row(src, reinterpret_cast(scratch.data()), + (int64_t)elems); + } + data = scratch.data(); + return true; +} + +size_t moe_expert_required_shared_bytes(int n_embd, + int n_expert_used, + int batch_limit, + ggml_type input_type) { + if (n_embd <= 0 || n_expert_used <= 0 || batch_limit <= 0) return 0; + size_t input_bytes = 0; + size_t ids_bytes = 0; + size_t weights_bytes = 0; + size_t total = 0; + size_t input_elems = 0; + size_t selected_elems = 0; + if (!moe_expert_compute_checked_mul_size((size_t)batch_limit, (size_t)n_embd, input_elems) || + !moe_expert_compute_checked_mul_size((size_t)batch_limit, (size_t)n_expert_used, selected_elems) || + !moe_expert_compute_checked_mul_size( + (size_t)batch_limit, + moe_expert_ipc_input_row_size(input_type, (size_t)n_embd), + input_bytes) || + !moe_expert_compute_checked_mul_size(selected_elems, sizeof(int32_t), ids_bytes) || + !moe_expert_compute_checked_mul_size(selected_elems, sizeof(float), weights_bytes) || + !backend_ipc_checked_add_size(input_bytes, ids_bytes, total) || + !backend_ipc_checked_add_size(total, weights_bytes, total)) { + return 0; + } + size_t response_bytes = 0; + if (!moe_expert_compute_checked_mul_size(input_elems, sizeof(float), + response_bytes)) { + return 0; + } + return std::max(total, response_bytes); +} + +size_t moe_expert_shared_bytes_from_env(size_t required_bytes) { + const char * raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_IPC_SHARED_BYTES"); + if (!raw || !*raw) return required_bytes; + if (raw[0] == '-') return required_bytes; + char * end = nullptr; + const unsigned long long parsed = std::strtoull(raw, &end, 10); + if (end == raw || *end != '\0' || + parsed > (unsigned long long)std::numeric_limits::max()) { + return required_bytes; + } + return std::max((size_t)parsed, required_bytes); +} + +bool open_expert_mmap(const char * target_path, + gguf_context *& gctx, + void *& mmap_addr, + size_t & file_size, + const uint8_t *& file_bytes, + size_t & data_start, + std::string * err) { +#if defined(_WIN32) + (void)target_path; (void)gctx; (void)mmap_addr; (void)file_size; + (void)file_bytes; (void)data_start; + if (err) *err = "mmap expert loading is POSIX-only"; + return false; +#else + gguf_init_params gip{}; + gctx = gguf_init_from_file(target_path, gip); + if (!gctx) { + if (err) *err = "failed to open GGUF for expert mmap"; + return false; + } + int fd = ::open(target_path, O_RDONLY); + if (fd < 0) { + if (err) *err = "failed to open target GGUF"; + gguf_free(gctx); + gctx = nullptr; + return false; + } + struct stat st; + if (::fstat(fd, &st) < 0) { + ::close(fd); + if (err) *err = "fstat failed on target GGUF"; + gguf_free(gctx); + gctx = nullptr; + return false; + } + file_size = (size_t)st.st_size; + mmap_addr = ::mmap(nullptr, file_size, PROT_READ, MAP_PRIVATE, fd, 0); + ::close(fd); + if (mmap_addr == MAP_FAILED) { + if (err) *err = "mmap failed on target GGUF"; + gguf_free(gctx); + gctx = nullptr; + mmap_addr = nullptr; + return false; + } + data_start = gguf_get_data_offset(gctx); + file_bytes = static_cast(mmap_addr); + return true; +#endif +} + +void close_expert_mmap(gguf_context * gctx, void * mmap_addr, size_t file_size) { +#if !defined(_WIN32) + if (mmap_addr && mmap_addr != MAP_FAILED) { + ::munmap(mmap_addr, file_size); + } +#else + (void)mmap_addr; (void)file_size; +#endif + if (gctx) gguf_free(gctx); +} + +std::vector make_layer_expert_file_data( + gguf_context * gctx, + const uint8_t * file_bytes, + size_t file_size, + size_t data_start, + int n_layer) { + std::vector layer_file_data((size_t)n_layer); + for (int il = 0; il < n_layer; ++il) { + char name[128]; + auto find_tensor_data = [&](const char * suffix) -> ExpertTensorFileData { + std::snprintf(name, sizeof(name), "blk.%d.%s.weight", il, suffix); + int64_t tid = gguf_find_tensor(gctx, name); + if (tid < 0) return {}; + const size_t off = data_start + gguf_get_tensor_offset(gctx, tid); + const size_t sz = gguf_get_tensor_size(gctx, tid); + if (off + sz > file_size) return {}; + return { file_bytes + off, sz }; + }; + layer_file_data[(size_t)il].gate_exps = find_tensor_data("ffn_gate_exps"); + layer_file_data[(size_t)il].up_exps = find_tensor_data("ffn_up_exps"); + layer_file_data[(size_t)il].down_exps = find_tensor_data("ffn_down_exps"); + layer_file_data[(size_t)il].gate_up_exps = find_tensor_data("ffn_gate_up_exps"); + } + return layer_file_data; +} + +std::string read_gguf_arch(const char * path, std::string * err) { + gguf_init_params gip{}; + gguf_context * gctx = gguf_init_from_file(path, gip); + if (!gctx) { + if (err) *err = "failed to open GGUF for arch detection"; + return {}; + } + int64_t arch_id = gguf_find_key(gctx, "general.architecture"); + if (arch_id < 0) { + if (err) *err = "missing general.architecture"; + gguf_free(gctx); + return {}; + } + const char * arch = gguf_get_val_str(gctx, arch_id); + std::string out = arch ? arch : ""; + gguf_free(gctx); + return out; +} + +template +bool build_remote_moe_runtime_from_weights( + ggml_backend_t backend, + const MoeHybridPlacement & placement, + const GgufTensorDataReader & expert_reader, + const std::string & arch, + WeightsT & weights, + FreeFn free_weights, + RemoteMoeRuntime & out, + std::string * err) { + out.arch = arch; + out.n_layer = weights.n_layer; + out.n_embd = weights.n_embd; + out.n_expert = weights.n_expert; + out.n_expert_used = weights.n_expert_used; + out.n_ff_exp = weights.n_ff_exp; + + MoeHybridConfig cfg = make_moe_hybrid_config(weights); + std::vector layer_descs((size_t)weights.n_layer); + out.layers.resize((size_t)weights.n_layer); + for (int il = 0; il < weights.n_layer; ++il) { + layer_descs[(size_t)il] = make_moe_layer_desc(weights.layers[(size_t)il]); + out.layers[(size_t)il].gate_scale = layer_descs[(size_t)il].ffn_gate_exps_s; + out.layers[(size_t)il].up_scale = layer_descs[(size_t)il].ffn_up_exps_s; + out.layers[(size_t)il].down_scale = layer_descs[(size_t)il].ffn_down_exps_s; + out.layers[(size_t)il].gate_up_scale = layer_descs[(size_t)il].ffn_gate_up_exps_s; + } + + const auto layer_file_data = + make_layer_expert_file_data(expert_reader, weights.n_layer); + std::fprintf(stderr, + "[moe-expert-compute-daemon] building remote hot storage arch=%s\n", + arch.c_str()); + if (!build_moe_hybrid_storage_from_file(cfg, backend, placement, layer_descs, + layer_file_data, out.hybrid, err, + /*cache_slots=*/0, + /*load_cold_tensors=*/false)) { + free_weights(weights); + return false; + } + free_weights(weights); + return true; +} + +bool load_remote_moe_runtime(const char * target_path, + ggml_backend_t backend, + const MoeHybridPlacement & placement, + const GgufTensorDataReader & expert_reader, + RemoteMoeRuntime & out, + std::string * err) { + std::string arch = read_gguf_arch(target_path, err); + if (arch.empty()) return false; + + TargetLoadPlan plan; + plan.skip_expert_tensors = true; + plan.load_output = false; + + if (arch == "qwen35moe") { + TargetWeights weights; + if (!load_target_gguf_partial(target_path, backend, plan, weights)) { + if (err) *err = dflash27b_last_error(); + free_target_weights(weights); + return false; + } + return build_remote_moe_runtime_from_weights( + backend, placement, expert_reader, arch, weights, + free_target_weights, out, err); + } + if (arch == "laguna") { + LagunaTargetWeights weights; + if (!load_target_gguf_laguna_partial(target_path, backend, plan, weights)) { + if (err) *err = dflash27b_last_error(); + free_laguna_target_weights(weights); + return false; + } + return build_remote_moe_runtime_from_weights( + backend, placement, expert_reader, arch, weights, + free_laguna_target_weights, out, err); + } + + if (err) *err = "unsupported MoE expert compute arch: " + arch; + return false; +} + +bool validate_shared_payload_request(const void * shared_payload, + const void * shared_payload_data, + size_t shared_payload_capacity, + size_t bytes, + uint64_t seq) { + const auto * header = + static_cast(shared_payload); + if (!shared_payload || !shared_payload_data || seq == 0 || + !backend_ipc_payload_in_bounds(0, bytes, shared_payload_capacity) || + header->sequence != seq || header->bytes != (uint64_t)bytes) { + return false; + } + return true; +} + +bool commit_shared_payload_response(void * shared_payload, + void * shared_payload_data, + size_t shared_payload_capacity, + size_t bytes, + uint64_t seq) { + if (!shared_payload || !shared_payload_data || seq == 0 || + !backend_ipc_payload_in_bounds(0, bytes, shared_payload_capacity)) { + return false; + } + auto * header = static_cast(shared_payload); + header->bytes = (uint64_t)bytes; + header->sequence = seq; + return true; +} + +int remote_hot_expert_count(const MoeHybridLayerStorage & storage) { + if (storage.gate_up_hot) return (int)storage.gate_up_hot->ne[2]; + if (storage.gate_hot) return (int)storage.gate_hot->ne[2]; + if (storage.down_hot) return (int)storage.down_hot->ne[2]; + return (int)storage.hot_expert_ids.size(); +} + +class MoeExpertComputeIpc : public MoeExpertCompute { +public: + explicit MoeExpertComputeIpc(int n_ff_max, bool enable_fallback = true) + : fallback_(enable_fallback ? make_cpu_moe_expert_compute(n_ff_max) : nullptr) {} + + ~MoeExpertComputeIpc() override { + if (profile_) { + stats_.maybe_print(true); + } + if (!placement_path_.empty()) { +#if !defined(_WIN32) + ::unlink(placement_path_.c_str()); +#endif + } + } + + bool start(const std::string & bin, + const std::string & target_path, + int target_gpu, + const MoeHybridPlacement & main_placement, + int n_embd, + int n_expert_used, + const std::string & work_dir, + bool required) { +#if defined(_WIN32) + (void)bin; (void)target_path; (void)target_gpu; (void)main_placement; + (void)n_embd; (void)n_expert_used; (void)work_dir; (void)required; + return false; +#else + std::string err; + if (!write_temp_remote_placement(main_placement, placement_path_, &err)) { + std::fprintf(stderr, "[moe-expert-compute-ipc] placement write failed: %s\n", + err.c_str()); + return false; + } + + BackendIpcLaunchConfig launch; + launch.bin = bin; + launch.mode = BackendIpcMode::MoeExpertCompute; + launch.payload_path = target_path; + launch.work_dir = work_dir; + launch.payload_transport = moe_expert_transport_from_env(); + input_type_ = moe_expert_ipc_input_type_from_env(); + const int shared_batch_capacity = + std::max(moe_expert_compute_batch_limit_from_env(), + moe_expert_ipc_shared_batch_capacity_from_env()); + launch.shared_payload_bytes = moe_expert_shared_bytes_from_env( + moe_expert_required_shared_bytes( + n_embd, n_expert_used, shared_batch_capacity, input_type_)); + launch.args.push_back("--target-gpu=" + std::to_string(std::max(0, target_gpu))); + launch.args.push_back("--placement=" + placement_path_); + if (!process_.start(launch)) { + std::fprintf(stderr, "[moe-expert-compute-ipc] backend process start failed%s\n", + required ? " (required)" : " (falling back to CPU)"); + return false; + } + active_ = true; + required_ = required; + n_embd_ = n_embd; + n_expert_used_ = n_expert_used; + profile_ = moe_expert_compute_profile_enabled(); + const size_t max_input_bytes = + moe_expert_ipc_input_row_size(input_type_, (size_t)n_embd_) * + (size_t)std::max(1, moe_expert_compute_batch_limit_from_env()); + if (max_input_bytes > input_scratch_.capacity()) { + input_scratch_.reserve(max_input_bytes); + } + std::printf("[moe-expert-compute-ipc] ready bin=%s gpu=%d input_type=%s work_dir=%s\n", + bin.c_str(), target_gpu, + ggml_type_name(input_type_), process_.work_dir().c_str()); + return true; +#endif + } + + bool start_for_placement(const std::string & bin, + const std::string & target_path, + int target_gpu, + const MoeHybridPlacement & remote_placement, + int n_embd, + int n_expert_used, + const std::string & work_dir, + bool required) { +#if defined(_WIN32) + (void)bin; (void)target_path; (void)target_gpu; (void)remote_placement; + (void)n_embd; (void)n_expert_used; (void)work_dir; (void)required; + return false; +#else + std::string err; + if (!write_temp_placement_file(remote_placement, "moe_remote_expert_target", + placement_path_, &err)) { + std::fprintf(stderr, "[moe-expert-compute-ipc] placement write failed: %s\n", + err.c_str()); + return false; + } + + BackendIpcLaunchConfig launch; + launch.bin = bin; + launch.mode = BackendIpcMode::MoeExpertCompute; + launch.payload_path = target_path; + launch.work_dir = work_dir; + launch.payload_transport = moe_expert_transport_from_env(); + input_type_ = moe_expert_ipc_input_type_from_env(); + const int shared_batch_capacity = + std::max(moe_expert_compute_batch_limit_from_env(), + moe_expert_ipc_shared_batch_capacity_from_env()); + launch.shared_payload_bytes = moe_expert_shared_bytes_from_env( + moe_expert_required_shared_bytes( + n_embd, n_expert_used, shared_batch_capacity, input_type_)); + launch.args.push_back("--target-gpu=" + std::to_string(std::max(0, target_gpu))); + launch.args.push_back("--placement=" + placement_path_); + if (!process_.start(launch)) { + std::fprintf(stderr, "[moe-expert-compute-ipc] backend process start failed%s\n", + required ? " (required)" : " (falling back to CPU)"); + return false; + } + active_ = true; + required_ = required; + n_embd_ = n_embd; + n_expert_used_ = n_expert_used; + profile_ = moe_expert_compute_profile_enabled(); + const size_t max_input_bytes = + moe_expert_ipc_input_row_size(input_type_, (size_t)n_embd_) * + (size_t)std::max(1, moe_expert_compute_batch_limit_from_env()); + if (max_input_bytes > input_scratch_.capacity()) { + input_scratch_.reserve(max_input_bytes); + } + std::printf("[moe-expert-compute-ipc] ready bin=%s gpu=%d input_type=%s work_dir=%s\n", + bin.c_str(), target_gpu, + ggml_type_name(input_type_), process_.work_dir().c_str()); + return true; +#endif + } + + bool compute(const MoeExpertLayer & layer, + const float * input, + const int32_t * ids, + const float * weights, + int n_cold, + int n_embd, + int n_ff, + float * output) override { + if (!output || n_embd <= 0) return false; + if (n_cold <= 0) { + std::fill(output, output + n_embd, 0.0f); + return true; + } + if (!active_ || layer.layer_idx < 0 || n_embd != n_embd_) { + if (required_) return false; + if (!fallback_) return false; + return fallback_->compute(layer, input, ids, weights, n_cold, n_embd, n_ff, output); + } + + std::fill(output, output + n_embd, 0.0f); + for (int i = 0; i < n_cold; ++i) { + const int32_t local = ids[i]; + if (local < 0 || (size_t)local >= layer.cold_global_by_local.size()) { + if (required_) return false; + if (!fallback_) return false; + return fallback_->compute(layer, input, ids, weights, n_cold, n_embd, n_ff, output); + } + } + + if (!compute_remote(layer.layer_idx, input, ids, weights, + n_cold, n_embd, output)) { + active_ = false; + if (required_) { + std::fprintf(stderr, + "[moe-expert-compute-ipc] required remote compute failed\n"); + return false; + } + if (!fallback_) return false; + return fallback_->compute(layer, input, ids, weights, n_cold, n_embd, n_ff, output); + } + return true; + } + + bool compute_batch(const MoeExpertLayer & layer, + const float * input, + const int32_t * ids, + const float * weights, + int n_tokens, + int n_selected, + int n_embd, + int n_ff, + float * output) override { + if (!output || n_tokens < 0 || n_selected < 0 || n_embd <= 0) { + return false; + } + if (n_tokens == 0 || n_selected == 0) { + std::fill(output, output + (size_t)n_tokens * (size_t)n_embd, 0.0f); + return true; + } + if (!active_ || layer.layer_idx < 0 || n_embd != n_embd_) { + if (required_) return false; + if (!fallback_) return false; + return fallback_->compute_batch(layer, input, ids, weights, + n_tokens, n_selected, + n_embd, n_ff, output); + } + + const int daemon_batch_limit = moe_expert_compute_daemon_batch_limit_from_env(); + if (n_tokens > daemon_batch_limit) { + for (int t0 = 0; t0 < n_tokens; t0 += daemon_batch_limit) { + const int tc = std::min(daemon_batch_limit, n_tokens - t0); + if (!compute_batch( + layer, + input + (size_t)t0 * (size_t)n_embd, + ids + (size_t)t0 * (size_t)n_selected, + weights + (size_t)t0 * (size_t)n_selected, + tc, + n_selected, + n_embd, + n_ff, + output + (size_t)t0 * (size_t)n_embd)) { + return false; + } + } + return true; + } + + for (int t = 0; t < n_tokens; ++t) { + for (int i = 0; i < n_selected; ++i) { + const size_t idx = (size_t)t * (size_t)n_selected + (size_t)i; + const int32_t local = ids[idx]; + if (local < 0 || (size_t)local >= layer.cold_global_by_local.size()) { + if (required_) return false; + if (!fallback_) return false; + return fallback_->compute_batch(layer, input, ids, weights, + n_tokens, n_selected, + n_embd, n_ff, output); + } + } + } + + if (!compute_remote_batch(layer.layer_idx, input, ids, + weights, n_tokens, + n_selected, n_embd, output)) { + active_ = false; + if (required_) { + std::fprintf(stderr, + "[moe-expert-compute-ipc] required remote batch compute failed\n"); + return false; + } + if (!fallback_) return false; + return fallback_->compute_batch(layer, input, ids, weights, + n_tokens, n_selected, + n_embd, n_ff, output); + } + return true; + } + + bool prepare_batch(const MoeExpertLayer & layer, + int n_tokens, + int n_selected, + int n_embd, + int n_ff) override { + (void)n_ff; + if (n_tokens < 0 || n_selected < 0 || n_embd <= 0) return false; + if (n_tokens == 0 || n_selected == 0) return true; + if (!active_ || layer.layer_idx < 0 || n_embd != n_embd_) { + return !required_ && fallback_ != nullptr; + } + if (n_selected > n_expert_used_ || + n_selected > (int)layer.cold_global_by_local.size()) { + return !required_ && fallback_ != nullptr; + } + return prepare_remote_batch(layer.layer_idx, n_tokens, n_selected); + } + + bool prepare_single(const MoeExpertLayer & layer, + int n_selected, + int n_embd, + int n_ff) override { + (void)n_ff; + if (n_selected < 0 || n_embd <= 0) return false; + if (n_selected == 0) return true; + if (!active_ || layer.layer_idx < 0 || n_embd != n_embd_) { + return !required_ && fallback_ != nullptr; + } + if (n_selected > n_expert_used_ || + n_selected > (int)layer.cold_global_by_local.size()) { + return !required_ && fallback_ != nullptr; + } + return prepare_remote_single(layer.layer_idx, n_selected); + } + + bool healthy() const override { + return active_; + } + + bool prefers_padded_batch() const override { + return active_; + } + +private: + bool compute_remote(int layer_idx, + const float * input, + const int32_t * local_ids, + const float * weights, + int n_selected, + int n_embd, + float * output) { +#if defined(_WIN32) + (void)layer_idx; (void)input; (void)local_ids; (void)weights; + (void)n_selected; (void)n_embd; (void)output; + return false; +#else + std::lock_guard lock(command_mu_); + FILE * cmd = process_.command_stream(); + const int stream_fd = process_.stream_fd(); + const int payload_fd = process_.payload_fd(); + if (!cmd || stream_fd < 0 || layer_idx < 0 || !input || + !local_ids || !weights || n_selected <= 0 || + n_selected > n_expert_used_) { + return false; + } + + const ggml_type request_input_type = GGML_TYPE_F32; + const void * ipc_input = nullptr; + size_t input_bytes = 0; + if (!moe_expert_convert_input_to_ipc_type( + request_input_type, input, 1, n_embd, input_scratch_, + ipc_input, input_bytes)) { + return false; + } + const size_t ids_bytes = (size_t)n_selected * sizeof(int32_t); + const size_t weights_bytes = (size_t)n_selected * sizeof(float); + size_t payload_bytes = 0; + if (!backend_ipc_checked_add_size(input_bytes, ids_bytes, payload_bytes) || + !backend_ipc_checked_add_size(payload_bytes, weights_bytes, payload_bytes)) { + return false; + } + const bool use_shared = + process_.resolved_payload_transport() == BackendIpcPayloadTransport::Shared; + const auto pack_t0 = MoeExpertClock::now(); + uint64_t seq = 0; + const BackendIpcPayloadSegment segments[] = { + {ipc_input, input_bytes}, + {local_ids, ids_bytes}, + {weights, weights_bytes}, + }; + const auto pack_t1 = MoeExpertClock::now(); + + const auto send_t0 = MoeExpertClock::now(); + if (use_shared) { + if (!process_.write_shared_payload_segments(segments, 3, seq)) { + return false; + } + if (request_input_type == GGML_TYPE_F32) { + std::fprintf(cmd, "compute_local_shared %d %d %zu %" PRIu64 "\n", + layer_idx, n_selected, payload_bytes, seq); + } else { + std::fprintf(cmd, "compute_local_shared_typed %d %d %d %zu %" PRIu64 "\n", + layer_idx, n_selected, (int)request_input_type, + payload_bytes, seq); + } + std::fflush(cmd); + } else if (payload_fd >= 0) { + if (request_input_type == GGML_TYPE_F32) { + std::fprintf(cmd, "compute_local_pipe %d %d %zu\n", + layer_idx, n_selected, payload_bytes); + } else { + std::fprintf(cmd, "compute_local_pipe_typed %d %d %d %zu\n", + layer_idx, n_selected, (int)request_input_type, + payload_bytes); + } + std::fflush(cmd); + if (!write_exact_fd(payload_fd, ipc_input, input_bytes) || + !write_exact_fd(payload_fd, local_ids, ids_bytes) || + !write_exact_fd(payload_fd, weights, weights_bytes)) { + return false; + } + } else { + return false; + } + const auto send_t1 = MoeExpertClock::now(); + + int32_t status = -1; + const auto wait_t0 = MoeExpertClock::now(); + if (!read_exact_fd(stream_fd, &status, sizeof(status)) || status != 0) { + std::fprintf(stderr, "[moe-expert-compute-ipc] compute failed status=%d\n", status); + return false; + } + const auto wait_t1 = MoeExpertClock::now(); + const size_t output_bytes = sizeof(float) * (size_t)n_embd; + const bool ok = use_shared + ? process_.read_shared_payload(output, output_bytes, seq) + : read_exact_fd(stream_fd, output, output_bytes); + const auto recv_t1 = MoeExpertClock::now(); + if (profile_) { + stats_.calls++; + stats_.payload_bytes += payload_bytes; + stats_.response_bytes += output_bytes; + stats_.pack_us += moe_expert_compute_elapsed_us(pack_t0, pack_t1); + stats_.send_us += moe_expert_compute_elapsed_us(send_t0, send_t1); + stats_.wait_us += moe_expert_compute_elapsed_us(wait_t0, wait_t1); + stats_.recv_us += moe_expert_compute_elapsed_us(wait_t1, recv_t1); + stats_.maybe_print(false); + } + return ok; +#endif + } + + bool prepare_remote_batch(int layer_idx, int n_tokens, int n_selected) { +#if defined(_WIN32) + (void)layer_idx; (void)n_tokens; (void)n_selected; + return false; +#else + std::lock_guard lock(command_mu_); + FILE * cmd = process_.command_stream(); + const int stream_fd = process_.stream_fd(); + if (!cmd || stream_fd < 0 || layer_idx < 0 || + n_tokens <= 0 || n_selected <= 0 || + n_selected > n_expert_used_) { + return false; + } + if (input_type_ == GGML_TYPE_F32) { + std::fprintf(cmd, "prepare_batch_local %d %d %d\n", + layer_idx, n_tokens, n_selected); + } else { + std::fprintf(cmd, "prepare_batch_local_typed %d %d %d %d\n", + layer_idx, n_tokens, n_selected, (int)input_type_); + } + std::fflush(cmd); + + int32_t status = -1; + if (!read_exact_fd(stream_fd, &status, sizeof(status)) || status != 0) { + std::fprintf(stderr, + "[moe-expert-compute-ipc] prepare_batch failed status=%d\n", + status); + return false; + } + return true; +#endif + } + + bool prepare_remote_single(int layer_idx, int n_selected) { +#if defined(_WIN32) + (void)layer_idx; (void)n_selected; + return false; +#else + std::lock_guard lock(command_mu_); + FILE * cmd = process_.command_stream(); + const int stream_fd = process_.stream_fd(); + if (!cmd || stream_fd < 0 || layer_idx < 0 || + n_selected <= 0 || n_selected > n_expert_used_) { + return false; + } + // Decode-path compute currently consumes F32 hidden states, so the + // single-token warmup must build the same F32 CachedFfnGraph. + std::fprintf(cmd, "prepare_single_local %d %d\n", + layer_idx, n_selected); + std::fflush(cmd); + int32_t status = -1; + if (!read_exact_fd(stream_fd, &status, sizeof(status)) || status != 0) { + std::fprintf(stderr, + "[moe-expert-compute-ipc] prepare_single failed status=%d\n", + status); + return false; + } + return true; +#endif + } + + bool compute_remote_batch(int layer_idx, + const float * input, + const int32_t * local_ids, + const float * weights, + int n_tokens, + int n_selected, + int n_embd, + float * output) { +#if defined(_WIN32) + (void)layer_idx; (void)input; (void)local_ids; (void)weights; + (void)n_tokens; (void)n_selected; (void)n_embd; (void)output; + return false; +#else + std::lock_guard lock(command_mu_); + FILE * cmd = process_.command_stream(); + const int stream_fd = process_.stream_fd(); + const int payload_fd = process_.payload_fd(); + if (!cmd || stream_fd < 0 || layer_idx < 0 || !input || + !local_ids || !weights || n_tokens <= 0 || n_selected <= 0 || + n_selected > n_expert_used_) { + return false; + } + + const ggml_type request_input_type = input_type_; + const void * ipc_input = nullptr; + size_t input_bytes = 0; + if (!moe_expert_convert_input_to_ipc_type( + request_input_type, input, n_tokens, n_embd, input_scratch_, + ipc_input, input_bytes)) { + return false; + } + const size_t ids_bytes = + (size_t)n_tokens * (size_t)n_selected * sizeof(int32_t); + const size_t weights_bytes = + (size_t)n_tokens * (size_t)n_selected * sizeof(float); + size_t payload_bytes = 0; + if (!backend_ipc_checked_add_size(input_bytes, ids_bytes, payload_bytes) || + !backend_ipc_checked_add_size(payload_bytes, weights_bytes, payload_bytes)) { + return false; + } + const bool use_shared = + process_.resolved_payload_transport() == BackendIpcPayloadTransport::Shared; + const auto pack_t0 = MoeExpertClock::now(); + uint64_t seq = 0; + const BackendIpcPayloadSegment segments[] = { + {ipc_input, input_bytes}, + {local_ids, ids_bytes}, + {weights, weights_bytes}, + }; + const auto pack_t1 = MoeExpertClock::now(); + + const auto send_t0 = MoeExpertClock::now(); + if (use_shared) { + if (!process_.write_shared_payload_segments(segments, 3, seq)) { + return false; + } + if (request_input_type == GGML_TYPE_F32) { + std::fprintf(cmd, "compute_batch_local_shared %d %d %d %zu %" PRIu64 "\n", + layer_idx, n_tokens, n_selected, payload_bytes, seq); + } else { + std::fprintf(cmd, "compute_batch_local_shared_typed %d %d %d %d %zu %" PRIu64 "\n", + layer_idx, n_tokens, n_selected, (int)request_input_type, + payload_bytes, seq); + } + std::fflush(cmd); + } else if (payload_fd >= 0) { + if (request_input_type == GGML_TYPE_F32) { + std::fprintf(cmd, "compute_batch_local_pipe %d %d %d %zu\n", + layer_idx, n_tokens, n_selected, payload_bytes); + } else { + std::fprintf(cmd, "compute_batch_local_pipe_typed %d %d %d %d %zu\n", + layer_idx, n_tokens, n_selected, (int)request_input_type, + payload_bytes); + } + std::fflush(cmd); + if (!write_exact_fd(payload_fd, ipc_input, input_bytes) || + !write_exact_fd(payload_fd, local_ids, ids_bytes) || + !write_exact_fd(payload_fd, weights, weights_bytes)) { + return false; + } + } else { + return false; + } + const auto send_t1 = MoeExpertClock::now(); + + int32_t status = -1; + const auto wait_t0 = MoeExpertClock::now(); + if (!read_exact_fd(stream_fd, &status, sizeof(status)) || status != 0) { + std::fprintf(stderr, "[moe-expert-compute-ipc] batch compute failed status=%d\n", status); + return false; + } + const auto wait_t1 = MoeExpertClock::now(); + const size_t output_bytes = + (size_t)n_tokens * (size_t)n_embd * sizeof(float); + const bool ok = use_shared + ? process_.read_shared_payload(output, output_bytes, seq) + : read_exact_fd(stream_fd, output, output_bytes); + const auto recv_t1 = MoeExpertClock::now(); + if (profile_) { + stats_.calls++; + stats_.payload_bytes += payload_bytes; + stats_.response_bytes += output_bytes; + stats_.pack_us += moe_expert_compute_elapsed_us(pack_t0, pack_t1); + stats_.send_us += moe_expert_compute_elapsed_us(send_t0, send_t1); + stats_.wait_us += moe_expert_compute_elapsed_us(wait_t0, wait_t1); + stats_.recv_us += moe_expert_compute_elapsed_us(wait_t1, recv_t1); + stats_.maybe_print(false); + } + return ok; +#endif + } + + BackendIpcProcess process_; + std::unique_ptr fallback_; + std::string placement_path_; + std::vector input_scratch_; + std::mutex command_mu_; + MoeExpertIpcClientStats stats_; + int n_embd_ = 0; + int n_expert_used_ = 0; + ggml_type input_type_ = GGML_TYPE_F32; + bool active_ = false; + bool required_ = false; + bool profile_ = false; +}; + +} // namespace + +MoeExpertComputeIpcStartResult make_moe_expert_compute_ipc( + const std::string & bin, + const std::string & target_path, + int target_gpu, + const MoeHybridPlacement & main_placement, + int n_embd, + int n_ff_exp, + int n_expert_used, + const std::string & work_dir, + bool required) { + MoeExpertComputeIpcStartResult result; + auto compute = std::make_unique(n_ff_exp); + if (!compute->start(bin, target_path, target_gpu, main_placement, + n_embd, n_expert_used, work_dir, required)) { + if (required) return result; + result.compute = make_cpu_moe_expert_compute(n_ff_exp); + return result; + } + result.started_remote = true; + result.compute = std::move(compute); + return result; +} + +MoeExpertComputeIpcStartResult make_moe_expert_compute_ipc_for_placement( + const std::string & bin, + const std::string & target_path, + int target_gpu, + const MoeHybridPlacement & remote_placement, + int n_embd, + int n_ff_exp, + int n_expert_used, + const std::string & work_dir, + bool required) { + MoeExpertComputeIpcStartResult result; + auto compute = std::make_unique( + n_ff_exp, /*enable_fallback=*/false); + if (!compute->start_for_placement(bin, target_path, target_gpu, remote_placement, + n_embd, n_expert_used, work_dir, required)) { + return result; + } + result.started_remote = true; + result.compute = std::move(compute); + return result; +} + +int run_moe_expert_compute_ipc_daemon(const char * target_path, + const char * placement_path, + int target_gpu, + int stream_fd, + int payload_fd, + int shared_payload_fd, + size_t shared_payload_bytes) { +#if defined(_WIN32) + (void)target_path; (void)placement_path; (void)target_gpu; + (void)stream_fd; (void)payload_fd; (void)shared_payload_fd; + (void)shared_payload_bytes; + std::fprintf(stderr, "MoE expert compute IPC daemon is only implemented on POSIX hosts\n"); + return 2; +#else + setvbuf(stderr, nullptr, _IONBF, 0); + if (!target_path || !placement_path || stream_fd < 0) { + std::fprintf(stderr, + "usage: backend_ipc_daemon --backend-ipc-mode=moe-expert-compute " + " --target-gpu=N --placement=PATH --stream-fd=FD\n"); + return 2; + } + + void * shared_payload = nullptr; + void * shared_payload_data = nullptr; + size_t shared_payload_capacity = 0; + size_t shared_payload_map_bytes = 0; + if (shared_payload_fd >= 0 || shared_payload_bytes > 0) { + if (shared_payload_fd < 0 || shared_payload_bytes == 0 || + !backend_ipc_shared_payload_map_bytes(shared_payload_bytes, + shared_payload_map_bytes)) { + emit_status(stream_fd, -1); + return 1; + } + shared_payload = ::mmap(nullptr, shared_payload_map_bytes, + PROT_READ | PROT_WRITE, MAP_SHARED, + shared_payload_fd, 0); + if (shared_payload == MAP_FAILED) { + std::fprintf(stderr, "[moe-expert-compute-daemon] shared payload mmap failed\n"); + emit_status(stream_fd, -1); + return 1; + } + shared_payload_data = + static_cast(shared_payload) + backend_ipc_shared_payload_header_bytes(); + shared_payload_capacity = shared_payload_bytes; + } + + ggml_backend_t backend = ggml_backend_cuda_init(std::max(0, target_gpu)); + if (!backend) { + std::fprintf(stderr, "[moe-expert-compute-daemon] backend init failed gpu=%d\n", + target_gpu); + emit_status(stream_fd, -1); + if (shared_payload && shared_payload != MAP_FAILED) { + ::munmap(shared_payload, shared_payload_map_bytes); + } + return 1; + } + + std::fprintf(stderr, + "[moe-expert-compute-daemon] starting target metadata load target=%s gpu=%d\n", + target_path, target_gpu); + + MoeHybridPlacement placement; + std::string err; + std::fprintf(stderr, "[moe-expert-compute-daemon] loading placement=%s\n", placement_path); + if (!MoeHybridPlacement::load_json(placement_path, placement, &err)) { + std::fprintf(stderr, "[moe-expert-compute-daemon] placement load failed: %s\n", + err.c_str()); + emit_status(stream_fd, -1); + ggml_backend_free(backend); + if (shared_payload && shared_payload != MAP_FAILED) { + ::munmap(shared_payload, shared_payload_map_bytes); + } + return 1; + } + std::fprintf(stderr, + "[moe-expert-compute-daemon] placement loaded total_hot=%d\n", + placement.total_hot); + + GgufTensorDataReader expert_reader; + std::fprintf(stderr, "[moe-expert-compute-daemon] opening expert GGUF shards\n"); + if (!expert_reader.open(target_path, + /*build_merged_tensor_context=*/false, err) || + !expert_reader.open_mmaps(err)) { + std::fprintf(stderr, "[moe-expert-compute-daemon] expert GGUF open failed: %s\n", + err.c_str()); + emit_status(stream_fd, -1); + ggml_backend_free(backend); + if (shared_payload && shared_payload != MAP_FAILED) { + ::munmap(shared_payload, shared_payload_map_bytes); + } + return 1; + } + std::fprintf(stderr, + "[moe-expert-compute-daemon] expert GGUF shards ready count=%d\n", + expert_reader.shard_count()); + + RemoteMoeRuntime runtime; + if (!load_remote_moe_runtime(target_path, backend, placement, + expert_reader, runtime, &err)) { + std::fprintf(stderr, "[moe-expert-compute-daemon] target runtime load failed: %s\n", + err.c_str()); + emit_status(stream_fd, -1); + ggml_backend_free(backend); + if (shared_payload && shared_payload != MAP_FAILED) { + ::munmap(shared_payload, shared_payload_map_bytes); + } + return 1; + } + std::fprintf(stderr, + "[moe-expert-compute-daemon] target runtime ready arch=%s layers=%d embd=%d experts=%d used=%d remote_hot=%d\n", + runtime.arch.c_str(), runtime.n_layer, runtime.n_embd, + runtime.n_expert, runtime.n_expert_used, + runtime.hybrid.placement.total_hot); + + std::vector> graphs( + (size_t)runtime.n_layer, + std::vector((size_t)runtime.n_expert_used + 1)); + std::vector> batch_graphs( + (size_t)runtime.n_layer); + std::vector input; + std::vector local_ids((size_t)std::max(1, runtime.n_expert_used)); + std::vector router_weights((size_t)std::max(1, runtime.n_expert_used)); + std::vector output((size_t)runtime.n_embd); + std::vector batch_local_ids; + std::vector payload_ids; + std::vector payload_weights; + MoeExpertDaemonStats profile_stats; + MoeExpertDaemonStats profile_stats_prefill; + MoeExpertDaemonStats profile_stats_decode; + const bool profile = moe_expert_compute_profile_enabled(); + const int reserve_batch_limit = std::max(1, moe_expert_compute_daemon_batch_limit_from_env()); + input.reserve((size_t)runtime.n_embd * (size_t)reserve_batch_limit * + sizeof(float)); + output.reserve((size_t)runtime.n_embd * (size_t)reserve_batch_limit); + batch_local_ids.reserve((size_t)reserve_batch_limit * (size_t)runtime.n_expert_used); + payload_ids.reserve((size_t)reserve_batch_limit * (size_t)runtime.n_expert_used); + payload_weights.reserve((size_t)reserve_batch_limit * (size_t)runtime.n_expert_used); + auto record_graph_build = [&](MoeExpertDaemonStats & scope_stats, + uint64_t build_us) { + profile_stats.graph_builds++; + profile_stats.graph_build_us += build_us; + scope_stats.graph_builds++; + scope_stats.graph_build_us += build_us; + }; + auto record_graph_reuse = [&](MoeExpertDaemonStats & scope_stats) { + profile_stats.graph_reuses++; + scope_stats.graph_reuses++; + }; + int warmup_builds = 0; + std::fprintf(stderr, + "[moe-expert-compute-daemon] ready gpu=%d remote_hot=%d layers=%d warmup_graphs=%d warmup_ms=0.000\n", + target_gpu, runtime.hybrid.placement.total_hot, runtime.n_layer, warmup_builds); + emit_status(stream_fd, 0); + + std::string line; + while (std::getline(std::cin, line)) { + std::istringstream iss(line); + std::string cmd; + iss >> cmd; + if (cmd == "quit" || cmd == "exit") { + break; + } + + int layer_idx = -1; + int n_selected = 0; + int n_tokens = 1; + size_t bytes = 0; + uint64_t shared_seq = 0; + bool payload_ok = false; + bool pipe_payload = false; + bool shared_payload_cmd = false; + bool ids_are_local = false; + bool prepare_only = false; + bool prepare_single_only = false; + ggml_type input_type = GGML_TYPE_F32; + const auto read_t0 = MoeExpertClock::now(); + if (cmd == "compute_pipe") { + iss >> layer_idx >> n_selected >> bytes; + pipe_payload = true; + } else if (cmd == "compute_shared") { + iss >> layer_idx >> n_selected >> bytes >> shared_seq; + shared_payload_cmd = true; + } else if (cmd == "compute_local_pipe") { + iss >> layer_idx >> n_selected >> bytes; + pipe_payload = true; + ids_are_local = true; + } else if (cmd == "compute_local_pipe_typed") { + int input_type_i = (int)GGML_TYPE_F32; + iss >> layer_idx >> n_selected >> input_type_i >> bytes; + pipe_payload = true; + ids_are_local = true; + input_type = (ggml_type)input_type_i; + } else if (cmd == "compute_local_shared") { + iss >> layer_idx >> n_selected >> bytes >> shared_seq; + shared_payload_cmd = true; + ids_are_local = true; + } else if (cmd == "compute_local_shared_typed") { + int input_type_i = (int)GGML_TYPE_F32; + iss >> layer_idx >> n_selected >> input_type_i >> bytes >> shared_seq; + shared_payload_cmd = true; + ids_are_local = true; + input_type = (ggml_type)input_type_i; + } else if (cmd == "compute_batch_pipe") { + iss >> layer_idx >> n_tokens >> n_selected >> bytes; + pipe_payload = true; + } else if (cmd == "compute_batch_shared") { + iss >> layer_idx >> n_tokens >> n_selected >> bytes >> shared_seq; + shared_payload_cmd = true; + } else if (cmd == "compute_batch_local_pipe") { + iss >> layer_idx >> n_tokens >> n_selected >> bytes; + pipe_payload = true; + ids_are_local = true; + } else if (cmd == "compute_batch_local_pipe_typed") { + int input_type_i = (int)GGML_TYPE_F32; + iss >> layer_idx >> n_tokens >> n_selected >> input_type_i >> bytes; + pipe_payload = true; + ids_are_local = true; + input_type = (ggml_type)input_type_i; + } else if (cmd == "compute_batch_local_shared") { + iss >> layer_idx >> n_tokens >> n_selected >> bytes >> shared_seq; + shared_payload_cmd = true; + ids_are_local = true; + } else if (cmd == "compute_batch_local_shared_typed") { + int input_type_i = (int)GGML_TYPE_F32; + iss >> layer_idx >> n_tokens >> n_selected >> input_type_i >> bytes >> shared_seq; + shared_payload_cmd = true; + ids_are_local = true; + input_type = (ggml_type)input_type_i; + } else if (cmd == "prepare_batch_local") { + iss >> layer_idx >> n_tokens >> n_selected; + ids_are_local = true; + prepare_only = true; + } else if (cmd == "prepare_batch_local_typed") { + int input_type_i = (int)GGML_TYPE_F32; + iss >> layer_idx >> n_tokens >> n_selected >> input_type_i; + ids_are_local = true; + prepare_only = true; + input_type = (ggml_type)input_type_i; + } else if (cmd == "prepare_single_local") { + iss >> layer_idx >> n_selected; + n_tokens = 1; + ids_are_local = true; + prepare_only = true; + prepare_single_only = true; + } else { + emit_status(stream_fd, -1); + continue; + } + + size_t input_elems = 0; + size_t selected_elems = 0; + size_t input_bytes = 0; + size_t ids_bytes = 0; + size_t weights_bytes = 0; + size_t expected = 0; + bool header_ok = false; + if (prepare_only) { + header_ok = n_tokens > 0 && n_selected > 0 && + n_selected <= runtime.n_expert_used && + layer_idx >= 0 && layer_idx < runtime.n_layer && + moe_expert_ipc_is_supported_input_type(input_type); + payload_ok = header_ok; + } else if (n_tokens > 0 && n_selected > 0 && + n_selected <= runtime.n_expert_used && + layer_idx >= 0 && layer_idx < runtime.n_layer && + moe_expert_ipc_is_supported_input_type(input_type)) { + header_ok = + moe_expert_compute_checked_mul_size((size_t)n_tokens, + (size_t)runtime.n_embd, + input_elems) && + moe_expert_compute_checked_mul_size((size_t)n_tokens, + (size_t)n_selected, + selected_elems) && + moe_expert_compute_checked_mul_size((size_t)n_tokens, + moe_expert_ipc_input_row_size( + input_type, + (size_t)runtime.n_embd), + input_bytes) && + moe_expert_compute_checked_mul_size(selected_elems, sizeof(int32_t), + ids_bytes) && + moe_expert_compute_checked_mul_size(selected_elems, sizeof(float), + weights_bytes) && + backend_ipc_checked_add_size(input_bytes, ids_bytes, expected) && + backend_ipc_checked_add_size(expected, weights_bytes, expected) && + bytes == expected; + } + + if (pipe_payload && header_ok) { + input.resize(input_bytes); + payload_ids.resize((size_t)n_tokens * (size_t)n_selected); + payload_weights.resize((size_t)n_tokens * (size_t)n_selected); + payload_ok = payload_fd >= 0 && + read_exact_fd(payload_fd, input.data(), input_bytes) && + read_exact_fd(payload_fd, payload_ids.data(), ids_bytes) && + read_exact_fd(payload_fd, payload_weights.data(), weights_bytes); + } else if (pipe_payload) { + payload_ok = payload_fd >= 0 && + moe_expert_compute_drain_exact_fd(payload_fd, bytes); + } else if (shared_payload_cmd) { + payload_ok = header_ok && + validate_shared_payload_request( + shared_payload, shared_payload_data, shared_payload_capacity, + bytes, shared_seq); + } + const auto read_t1 = MoeExpertClock::now(); + + if (!header_ok || !payload_ok) { + emit_status(stream_fd, -1); + continue; + } + + MoeHybridLayerStorage & storage = runtime.hybrid.layers[(size_t)layer_idx]; + const auto unpack_t1 = MoeExpertClock::now(); + const int remote_hot_count = remote_hot_expert_count(storage); + + auto & layer_graphs = graphs[(size_t)layer_idx]; + if ((size_t)n_selected >= layer_graphs.size()) { + emit_status(stream_fd, -1); + continue; + } + const RemoteMoeLayerRuntime & L = runtime.layers[(size_t)layer_idx]; + if (prepare_only) { + if (prepare_single_only && input_type == GGML_TYPE_F32) { + CachedFfnGraph & graph = layer_graphs[(size_t)n_selected]; + if (!graph.valid() || graph.n_hot != n_selected) { + const auto build_t0 = MoeExpertClock::now(); + if (!build_cached_cold_graph(graph, backend, + storage.gate_hot, storage.up_hot, + storage.down_hot, storage.gate_up_hot, + L.gate_scale, L.up_scale, + L.down_scale, L.gate_up_scale, + runtime.n_embd, runtime.n_ff_exp, + n_selected)) { + emit_status(stream_fd, -1); + continue; + } + const auto build_t1 = MoeExpertClock::now(); + if (profile) { + record_graph_build(profile_stats_decode, + moe_expert_compute_elapsed_us(build_t0, build_t1)); + } + } else if (profile) { + record_graph_reuse(profile_stats_decode); + } + emit_status(stream_fd, 0); + continue; + } + auto & layer_batch_graphs = batch_graphs[(size_t)layer_idx]; + const uint64_t graph_key = + moe_expert_compute_batch_graph_key(n_tokens, n_selected, + input_type); + auto inserted = layer_batch_graphs.try_emplace(graph_key); + CachedBatchedMoeExpertGraph & graph = inserted.first->second; + if (!graph.valid() || graph.n_tokens != n_tokens || + graph.n_selected != n_selected || + graph.input_type != input_type) { + const auto build_t0 = MoeExpertClock::now(); + if (!build_cached_batched_cold_graph( + graph, backend, + storage.gate_hot, storage.up_hot, + storage.down_hot, storage.gate_up_hot, + L.gate_scale, L.up_scale, + L.down_scale, L.gate_up_scale, + runtime.n_embd, runtime.n_ff_exp, + n_selected, n_tokens, input_type)) { + emit_status(stream_fd, -1); + continue; + } + const auto build_t1 = MoeExpertClock::now(); + if (profile) { + MoeExpertDaemonStats & scope_stats = + n_tokens > 1 ? profile_stats_prefill : profile_stats_decode; + record_graph_build(scope_stats, + moe_expert_compute_elapsed_us(build_t0, build_t1)); + } + } else if (profile) { + MoeExpertDaemonStats & scope_stats = + n_tokens > 1 ? profile_stats_prefill : profile_stats_decode; + record_graph_reuse(scope_stats); + } + emit_status(stream_fd, 0); + continue; + } + + size_t off = 0; + if (!shared_payload_cmd) { + output.resize((size_t)n_tokens * (size_t)runtime.n_embd); + } + const uint8_t * input_data = nullptr; + const int32_t * payload_id_data = nullptr; + const float * in_weights = nullptr; + if (pipe_payload) { + input_data = input.data(); + payload_id_data = payload_ids.data(); + in_weights = payload_weights.data(); + } else { + const auto * payload_data = + static_cast(shared_payload_data); + input_data = payload_data + off; + off += input_bytes; + payload_id_data = reinterpret_cast(payload_data + off); + off += ids_bytes; + in_weights = reinterpret_cast(payload_data + off); + } + + if (n_tokens > 1 || input_type != GGML_TYPE_F32) { + auto & layer_batch_graphs = batch_graphs[(size_t)layer_idx]; + const uint64_t graph_key = + moe_expert_compute_batch_graph_key(n_tokens, n_selected, + input_type); + auto inserted = layer_batch_graphs.try_emplace(graph_key); + CachedBatchedMoeExpertGraph & graph = inserted.first->second; + if (!graph.valid() || graph.n_tokens != n_tokens || + graph.n_selected != n_selected || + graph.input_type != input_type) { + const auto build_t0 = MoeExpertClock::now(); + if (!build_cached_batched_cold_graph( + graph, backend, + storage.gate_hot, storage.up_hot, + storage.down_hot, storage.gate_up_hot, + L.gate_scale, L.up_scale, + L.down_scale, L.gate_up_scale, + runtime.n_embd, runtime.n_ff_exp, + n_selected, n_tokens, input_type)) { + emit_status(stream_fd, -1); + continue; + } + const auto build_t1 = MoeExpertClock::now(); + if (profile) { + MoeExpertDaemonStats & scope_stats = + n_tokens > 1 ? profile_stats_prefill : profile_stats_decode; + record_graph_build(scope_stats, + moe_expert_compute_elapsed_us(build_t0, build_t1)); + } + } else if (profile) { + MoeExpertDaemonStats & scope_stats = + n_tokens > 1 ? profile_stats_prefill : profile_stats_decode; + record_graph_reuse(scope_stats); + } + + bool ids_ok = true; + const int32_t * graph_ids_data = payload_id_data; + if (!ids_are_local) { + batch_local_ids.resize((size_t)n_tokens * (size_t)n_selected); + graph_ids_data = batch_local_ids.data(); + } + for (int t = 0; t < n_tokens && ids_ok; ++t) { + for (int i = 0; i < n_selected; ++i) { + const size_t idx = (size_t)t * (size_t)n_selected + + (size_t)i; + const int32_t id = payload_id_data[idx]; + const int32_t local = ids_are_local + ? id + : ((id >= 0 && id < (int)storage.hot_local_by_global.size()) + ? storage.hot_local_by_global[(size_t)id] : -1); + if (local < 0 || local >= remote_hot_count) { + ids_ok = false; + break; + } + if (!ids_are_local) { + batch_local_ids[idx] = local; + } + } + } + if (!ids_ok) { + emit_status(stream_fd, -1); + continue; + } + + const auto set_one_t0 = MoeExpertClock::now(); + ggml_backend_tensor_set(graph.inp, input_data, 0, input_bytes); + ggml_backend_tensor_set(graph.ids, graph_ids_data, 0, ids_bytes); + ggml_backend_tensor_set(graph.weights, in_weights, 0, weights_bytes); + const auto set_one_t1 = MoeExpertClock::now(); + auto st = ggml_backend_graph_compute(backend, graph.gf); + const auto compute_one_t1 = MoeExpertClock::now(); + if (st != GGML_STATUS_SUCCESS) { + emit_status(stream_fd, -1); + continue; + } + const size_t output_bytes = + (size_t)n_tokens * (size_t)runtime.n_embd * sizeof(float); + void * response_data = shared_payload_cmd + ? shared_payload_data : static_cast(output.data()); + ggml_backend_tensor_get(graph.output, response_data, 0, + output_bytes); + const auto get_one_t1 = MoeExpertClock::now(); + + const int32_t status = 0; + if (shared_payload_cmd && + !commit_shared_payload_response(shared_payload, shared_payload_data, + shared_payload_capacity, + output_bytes, shared_seq)) { + emit_status(stream_fd, -1); + continue; + } + if (!write_exact_fd(stream_fd, &status, sizeof(status)) || + (!shared_payload_cmd && + !write_exact_fd(stream_fd, output.data(), output_bytes))) { + break; + } + const auto write_t1 = MoeExpertClock::now(); + if (profile) { + MoeExpertDaemonStats & scope_stats = + n_tokens > 1 ? profile_stats_prefill : profile_stats_decode; + const uint64_t read_us = + moe_expert_compute_elapsed_us(read_t0, read_t1); + const uint64_t unpack_us = + moe_expert_compute_elapsed_us(read_t1, unpack_t1); + const uint64_t set_us = + moe_expert_compute_elapsed_us(set_one_t0, set_one_t1); + const uint64_t compute_us = + moe_expert_compute_elapsed_us(set_one_t1, compute_one_t1); + const uint64_t get_us = + moe_expert_compute_elapsed_us(compute_one_t1, get_one_t1); + const uint64_t write_us = + moe_expert_compute_elapsed_us(get_one_t1, write_t1); + auto record_stats = [&](MoeExpertDaemonStats & stats) { + stats.calls++; + stats.payload_bytes += bytes; + stats.response_bytes += output_bytes; + stats.read_us += read_us; + stats.unpack_us += unpack_us; + stats.tensor_set_us += set_us; + stats.compute_us += compute_us; + stats.tensor_get_us += get_us; + stats.write_us += write_us; + }; + record_stats(profile_stats); + record_stats(scope_stats); + profile_stats.maybe_print(false); + scope_stats.maybe_print(n_tokens > 1, + n_tokens > 1 ? "[prefill]" : "[decode]"); + } + continue; + } + + CachedFfnGraph & graph = layer_graphs[(size_t)n_selected]; + if (!graph.valid() || graph.n_hot != n_selected) { + const auto build_t0 = MoeExpertClock::now(); + if (!build_cached_cold_graph(graph, backend, + storage.gate_hot, storage.up_hot, + storage.down_hot, storage.gate_up_hot, + L.gate_scale, L.up_scale, + L.down_scale, L.gate_up_scale, + runtime.n_embd, runtime.n_ff_exp, + n_selected)) { + emit_status(stream_fd, -1); + continue; + } + const auto build_t1 = MoeExpertClock::now(); + if (profile) { + record_graph_build(profile_stats_decode, + moe_expert_compute_elapsed_us(build_t0, build_t1)); + } + } else if (profile) { + record_graph_reuse(profile_stats_decode); + } + + bool compute_ok = true; + uint64_t set_us_accum = 0; + uint64_t compute_us_accum = 0; + uint64_t get_us_accum = 0; + const size_t one_input_bytes = (size_t)runtime.n_embd * sizeof(float); + const size_t one_ids_bytes = (size_t)n_selected * sizeof(int32_t); + const size_t one_weights_bytes = (size_t)n_selected * sizeof(float); + for (int t = 0; t < n_tokens; ++t) { + bool ids_ok = true; + const size_t token_off = (size_t)t * (size_t)n_selected; + const int32_t * token_ids = payload_id_data + token_off; + const float * token_weights = in_weights + token_off; + const int32_t * graph_ids_data = token_ids; + if (!ids_are_local) { + graph_ids_data = local_ids.data(); + } + for (int i = 0; i < n_selected; ++i) { + const int32_t id = token_ids[i]; + const int32_t local = ids_are_local + ? id + : ((id >= 0 && id < (int)storage.hot_local_by_global.size()) + ? storage.hot_local_by_global[(size_t)id] : -1); + if (local < 0 || local >= remote_hot_count) { + ids_ok = false; + break; + } + if (!ids_are_local) { + local_ids[(size_t)i] = local; + } + } + if (!ids_ok) { + compute_ok = false; + break; + } + const auto set_one_t0 = MoeExpertClock::now(); + ggml_backend_tensor_set(graph.inp, + static_cast(input_data) + + (size_t)t * one_input_bytes, + 0, one_input_bytes); + ggml_backend_tensor_set(graph.ids, graph_ids_data, 0, + one_ids_bytes); + ggml_backend_tensor_set(graph.weights, token_weights, 0, + one_weights_bytes); + const auto set_one_t1 = MoeExpertClock::now(); + auto st = ggml_backend_graph_compute(backend, graph.gf); + const auto compute_one_t1 = MoeExpertClock::now(); + if (st != GGML_STATUS_SUCCESS) { + compute_ok = false; + break; + } + void * response_data = shared_payload_cmd + ? static_cast( + static_cast(shared_payload_data) + + (size_t)t * one_input_bytes) + : static_cast( + output.data() + (size_t)t * (size_t)runtime.n_embd); + ggml_backend_tensor_get(graph.output, response_data, 0, + one_input_bytes); + const auto get_one_t1 = MoeExpertClock::now(); + set_us_accum += moe_expert_compute_elapsed_us(set_one_t0, set_one_t1); + compute_us_accum += moe_expert_compute_elapsed_us(set_one_t1, compute_one_t1); + get_us_accum += moe_expert_compute_elapsed_us(compute_one_t1, get_one_t1); + } + const auto set_t1 = MoeExpertClock::now(); + const auto get_t1 = set_t1; + if (!compute_ok) { + emit_status(stream_fd, -1); + continue; + } + const int32_t status = 0; + if (shared_payload_cmd && + !commit_shared_payload_response(shared_payload, shared_payload_data, + shared_payload_capacity, + input_bytes, shared_seq)) { + emit_status(stream_fd, -1); + continue; + } + if (!write_exact_fd(stream_fd, &status, sizeof(status)) || + (!shared_payload_cmd && + !write_exact_fd(stream_fd, output.data(), input_bytes))) { + break; + } + const auto write_t1 = MoeExpertClock::now(); + if (profile) { + const uint64_t read_us = + moe_expert_compute_elapsed_us(read_t0, read_t1); + const uint64_t unpack_us = + moe_expert_compute_elapsed_us(read_t1, unpack_t1); + const uint64_t write_us = + moe_expert_compute_elapsed_us(get_t1, write_t1); + auto record_stats = [&](MoeExpertDaemonStats & stats) { + stats.calls++; + stats.payload_bytes += bytes; + stats.response_bytes += input_bytes; + stats.read_us += read_us; + stats.unpack_us += unpack_us; + stats.tensor_set_us += set_us_accum; + stats.compute_us += compute_us_accum; + stats.tensor_get_us += get_us_accum; + stats.write_us += write_us; + }; + record_stats(profile_stats); + record_stats(profile_stats_decode); + profile_stats.maybe_print(false); + profile_stats_decode.maybe_print(false, "[decode]"); + } + } + + if (profile) { + profile_stats.maybe_print(true); + profile_stats_prefill.maybe_print(true, "[prefill]"); + profile_stats_decode.maybe_print(true, "[decode]"); + } + for (auto & per_layer : graphs) { + for (auto & graph : per_layer) graph.free(); + } + for (auto & per_layer : batch_graphs) { + for (auto & graph : per_layer) graph.second.free(); + } + ggml_backend_free(backend); + if (shared_payload && shared_payload != MAP_FAILED) { + ::munmap(shared_payload, shared_payload_map_bytes); + } + return 0; +#endif +} + +} // namespace dflash::common diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index e8bddebd2..806243eed 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -6,9 +6,14 @@ #include #include +#include +#include #include #include #include +#include +#include +#include namespace dflash::common { @@ -25,6 +30,391 @@ static uint64_t elapsed_us(HybridClock::time_point start, HybridClock::time_poin return (uint64_t) std::chrono::duration_cast(end - start).count(); } +struct MixedRoutingScratch { + std::vector hot_sel; + std::vector hot_wts; + std::vector cold_sel; + std::vector cold_wts; + std::vector lut; + std::vector vlut; + + void ensure(int total_slots, int n_expert) { + if (total_slots > 0) { + hot_sel.resize((size_t)total_slots); + hot_wts.resize((size_t)total_slots); + cold_sel.resize((size_t)total_slots); + cold_wts.resize((size_t)total_slots); + } + if (n_expert > 0) { + lut.resize((size_t)n_expert); + vlut.resize((size_t)n_expert); + } + } +}; + +static thread_local MixedRoutingScratch g_mixed_routing_scratch; + +struct RemoteColdBatchScratch { + std::vector cold_counts; + std::vector cold_sel; + std::vector cold_wts; + std::vector> token_groups; + std::vector active_groups; + std::vector group_input; + std::vector group_output; + std::vector group_ids; + std::vector group_wts; + + void ensure(int n_tokens, int n_used, int n_embd, int max_cold_selected) { + if (n_tokens > 0) { + cold_counts.resize((size_t)n_tokens); + } + if (n_tokens > 0 && n_used > 0) { + cold_sel.resize((size_t)n_tokens * (size_t)n_used); + cold_wts.resize((size_t)n_tokens * (size_t)n_used); + } + if (n_embd > 0 && n_tokens > 0) { + group_input.reserve((size_t)n_embd * (size_t)n_tokens); + group_output.reserve((size_t)n_embd * (size_t)n_tokens); + } + if (max_cold_selected > 0 && n_tokens > 0) { + group_ids.reserve((size_t)n_tokens * (size_t)max_cold_selected); + group_wts.reserve((size_t)n_tokens * (size_t)max_cold_selected); + } + } +}; + +static thread_local RemoteColdBatchScratch g_remote_cold_batch_scratch; + +struct HybridSplitScratch { + std::vector cold_partial; + std::vector hot_out; +}; + +static thread_local HybridSplitScratch g_hybrid_split_scratch; + +static bool prepare_unique_hot_local_slots( + const MoeHybridLayerStorage & storage, + const int32_t * selected_ids, + const float * selected_weights, + int n_tokens, + int n_used, + std::vector & out_ids, + std::vector & out_weights, + bool * out_has_hot, + bool * out_has_cold, + std::string * err) { + const int n_hot_init = std::max(1, storage.hot_active); + std::vector used((size_t) n_hot_init, 0); + bool has_hot = false; + bool has_cold = false; + + for (int t = 0; t < n_tokens; ++t) { + std::fill(used.begin(), used.end(), 0); + int next_dummy = 0; + for (int i = 0; i < n_used; ++i) { + const size_t idx = (size_t) t * (size_t) n_used + (size_t) i; + out_ids[idx] = -1; + out_weights[idx] = 0.0f; + + const int32_t gid = selected_ids[idx]; + if (gid < 0 || gid >= (int32_t) storage.hot_local_by_global.size()) { + continue; + } + const int32_t hot_local = storage.hot_local_by_global[(size_t) gid]; + if (hot_local >= 0) { + if (hot_local >= n_hot_init) { + if (err) *err = "hot local id exceeds initialized hot stack"; + return false; + } + if (used[(size_t) hot_local]) { + if (err) *err = "duplicate hot local id in routed selection"; + return false; + } + out_ids[idx] = hot_local; + out_weights[idx] = selected_weights[idx]; + used[(size_t) hot_local] = 1; + has_hot = true; + continue; + } + if ((size_t) gid < storage.cold_local_by_global.size() && + storage.cold_local_by_global[(size_t) gid] >= 0) { + has_cold = true; + } + } + + for (int i = 0; i < n_used; ++i) { + const size_t idx = (size_t) t * (size_t) n_used + (size_t) i; + if (out_ids[idx] >= 0) continue; + while (next_dummy < n_hot_init && used[(size_t) next_dummy]) { + ++next_dummy; + } + if (next_dummy >= n_hot_init) { + if (err) *err = "insufficient initialized hot experts for unique padded routing"; + return false; + } + out_ids[idx] = next_dummy; + used[(size_t) next_dummy] = 1; + ++next_dummy; + } + } + + if (out_has_hot) *out_has_hot = has_hot; + if (out_has_cold) *out_has_cold = has_cold; + return true; +} + +static bool prepare_unique_cold_local_slots( + const MoeHybridLayerStorage & storage, + const int32_t * selected_ids, + const float * selected_weights, + int n_tokens, + int n_used, + std::vector & out_ids, + std::vector & out_weights, + bool * out_has_cold, + std::string * err) { + const int n_cold_stack = + std::max(1, (int) (storage.down_cold ? storage.down_cold->ne[2] : 0)); + std::vector used((size_t) n_cold_stack, 0); + bool has_cold = false; + + for (int t = 0; t < n_tokens; ++t) { + std::fill(used.begin(), used.end(), 0); + int next_dummy = 0; + for (int i = 0; i < n_used; ++i) { + const size_t idx = (size_t) t * (size_t) n_used + (size_t) i; + out_ids[idx] = -1; + out_weights[idx] = 0.0f; + + const int32_t gid = selected_ids[idx]; + if (gid < 0 || gid >= (int32_t) storage.hot_local_by_global.size()) { + continue; + } + if (storage.hot_local_by_global[(size_t) gid] >= 0) { + continue; + } + const int32_t cold_local = storage.cold_local_by_global[(size_t) gid]; + if (cold_local < 0) { + continue; + } + if (cold_local >= n_cold_stack) { + if (err) *err = "cold local id exceeds initialized cold stack"; + return false; + } + if (used[(size_t) cold_local]) { + if (err) *err = "duplicate cold local id in routed selection"; + return false; + } + out_ids[idx] = cold_local; + out_weights[idx] = selected_weights[idx]; + used[(size_t) cold_local] = 1; + has_cold = true; + } + + for (int i = 0; i < n_used; ++i) { + const size_t idx = (size_t) t * (size_t) n_used + (size_t) i; + if (out_ids[idx] >= 0) continue; + while (next_dummy < n_cold_stack && used[(size_t) next_dummy]) { + ++next_dummy; + } + if (next_dummy >= n_cold_stack) { + if (err) *err = "insufficient cold experts for unique padded routing"; + return false; + } + out_ids[idx] = next_dummy; + used[(size_t) next_dummy] = 1; + ++next_dummy; + } + } + + if (out_has_cold) *out_has_cold = has_cold; + return true; +} + +static bool prepare_unique_hot_global_slots( + const MoeHybridLayerStorage & storage, + const int32_t * selected_ids, + const float * selected_weights, + int n_tokens, + int n_used, + std::vector & out_global_ids, + std::vector & out_weights, + std::string * err) { + const int n_hot_init = std::max(1, storage.hot_active); + std::vector hot_global_by_local((size_t) n_hot_init, -1); + for (size_t global = 0; global < storage.hot_local_by_global.size(); ++global) { + const int32_t local = storage.hot_local_by_global[global]; + if (local >= 0 && local < n_hot_init) { + hot_global_by_local[(size_t) local] = (int32_t) global; + } + } + + std::vector used((size_t) n_hot_init, 0); + for (int t = 0; t < n_tokens; ++t) { + std::fill(used.begin(), used.end(), 0); + int next_dummy = 0; + for (int i = 0; i < n_used; ++i) { + const size_t idx = (size_t) t * (size_t) n_used + (size_t) i; + out_global_ids[idx] = -1; + out_weights[idx] = 0.0f; + + const int32_t gid = selected_ids[idx]; + if (gid < 0 || gid >= (int32_t) storage.hot_local_by_global.size()) { + continue; + } + const int32_t hot_local = storage.hot_local_by_global[(size_t) gid]; + if (hot_local < 0) { + continue; + } + if (hot_local >= n_hot_init) { + if (err) *err = "hot local id exceeds initialized hot stack"; + return false; + } + if (used[(size_t) hot_local]) { + if (err) *err = "duplicate hot local id in routed selection"; + return false; + } + out_global_ids[idx] = gid; + out_weights[idx] = selected_weights[idx]; + used[(size_t) hot_local] = 1; + } + + for (int i = 0; i < n_used; ++i) { + const size_t idx = (size_t) t * (size_t) n_used + (size_t) i; + if (out_global_ids[idx] >= 0) continue; + while (next_dummy < n_hot_init && + (used[(size_t) next_dummy] || + hot_global_by_local[(size_t) next_dummy] < 0)) { + ++next_dummy; + } + if (next_dummy >= n_hot_init) { + if (err) *err = "insufficient initialized hot experts for unique remapped routing"; + return false; + } + out_global_ids[idx] = hot_global_by_local[(size_t) next_dummy]; + used[(size_t) next_dummy] = 1; + ++next_dummy; + } + } + + return true; +} + +class HybridColdWorker { +public: + HybridColdWorker() = default; + ~HybridColdWorker() { + { + std::lock_guard lock(mu_); + stopping_ = true; + } + task_cv_.notify_one(); + if (worker_.joinable()) { + worker_.join(); + } + } + + HybridColdWorker(const HybridColdWorker &) = delete; + HybridColdWorker & operator=(const HybridColdWorker &) = delete; + + void submit(std::function task) { + ensure_started(); + { + std::unique_lock lock(mu_); + done_cv_.wait(lock, [this]() { return done_ && !has_task_; }); + task_ = std::move(task); + has_task_ = true; + done_ = false; + } + task_cv_.notify_one(); + } + + void wait() { + std::unique_lock lock(mu_); + done_cv_.wait(lock, [this]() { return done_; }); + } + +private: + void ensure_started() { + if (worker_.joinable()) return; + worker_ = std::thread([this]() { loop(); }); + } + + void loop() { + for (;;) { + std::function task; + { + std::unique_lock lock(mu_); + task_cv_.wait(lock, [this]() { return stopping_ || has_task_; }); + if (stopping_ && !has_task_) { + return; + } + task = std::move(task_); + has_task_ = false; + } + + task(); + + { + std::lock_guard lock(mu_); + done_ = true; + } + done_cv_.notify_one(); + } + } + + std::thread worker_; + std::mutex mu_; + std::condition_variable task_cv_; + std::condition_variable done_cv_; + std::function task_; + bool has_task_ = false; + bool done_ = true; + bool stopping_ = false; +}; + +static thread_local HybridColdWorker g_hybrid_cold_worker; + +static bool moe_hybrid_batched_profile_enabled() { + const char * raw = std::getenv("DFLASH_MOE_BATCHED_PROFILE"); + return raw && *raw && std::strcmp(raw, "0") != 0 && + std::strcmp(raw, "false") != 0 && std::strcmp(raw, "off") != 0; +} + +struct MoeHybridBatchedProfile { + uint64_t calls = 0; + uint64_t tokens = 0; + uint64_t partition_us = 0; + uint64_t remote_cold_us = 0; + uint64_t overlap_us = 0; + uint64_t hot_calls = 0; + uint64_t hot_tokens = 0; + uint64_t hot_us = 0; + uint64_t merge_us = 0; + uint64_t total_us = 0; + + void maybe_print(bool force) const { + if (calls == 0) return; + if (!force && calls % 64 != 0) return; + std::fprintf(stderr, + "[hybrid-ffn-batched] profile calls=%" PRIu64 + " tokens=%" PRIu64 " hot_calls=%" PRIu64 + " hot_tokens=%" PRIu64 + " partition_ms=%.3f remote_cold_ms=%.3f hot_ms=%.3f" + " overlap_ms=%.3f merge_ms=%.3f total_ms=%.3f\n", + calls, tokens, hot_calls, hot_tokens, + partition_us / 1000.0, remote_cold_us / 1000.0, + hot_us / 1000.0, overlap_us / 1000.0, + merge_us / 1000.0, total_us / 1000.0); + } +}; + +static MoeHybridBatchedProfile & moe_hybrid_batched_profile() { + static MoeHybridBatchedProfile profile; + return profile; +} + // Build the shared-expert FFN subgraph onto an existing ggml_context. // Returns the output tensor (or nullptr if no shared expert is present). static ggml_tensor * build_shared_expert_subgraph( @@ -617,10 +1007,12 @@ bool build_cached_hot_batched_graph( MoeHybridLayerStorage & storage, const MoeLayerDesc & desc, const MoeHybridConfig & cfg, - int n_tokens) { + int n_tokens, + bool gpu_remap) { out.free(); out.n_tokens = n_tokens; + out.gpu_remap = gpu_remap; const int n_embd = cfg.n_embd; const int n_used = cfg.n_expert_used; @@ -640,6 +1032,26 @@ bool build_cached_hot_batched_graph( out.wts = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, n_used, n_tokens); ggml_set_input(out.wts); + if (gpu_remap) { + const int total_slots = n_used * n_tokens; + out.global_ids = ggml_new_tensor_2d(out.ctx, GGML_TYPE_I32, total_slots, 1); + ggml_set_input(out.global_ids); + out.raw_weights = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, total_slots, 1); + ggml_set_input(out.raw_weights); + out.hot_local_lut = ggml_new_tensor_2d(out.ctx, GGML_TYPE_I32, 1, cfg.n_expert); + ggml_set_input(out.hot_local_lut); + out.valid_lut = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, 1, cfg.n_expert); + ggml_set_input(out.valid_lut); + + ggml_tensor * lid = ggml_get_rows(out.ctx, out.hot_local_lut, out.global_ids); + out.sel = ggml_cont(out.ctx, ggml_reshape_2d(out.ctx, lid, n_used, n_tokens)); + ggml_tensor * vm = ggml_get_rows(out.ctx, out.valid_lut, out.global_ids); + vm = ggml_reshape_2d(out.ctx, vm, n_used, n_tokens); + ggml_tensor * raw_wts = ggml_reshape_2d(out.ctx, out.raw_weights, + n_used, n_tokens); + out.wts = ggml_mul(out.ctx, raw_wts, vm); + } + ggml_tensor * routed = nullptr; build_batched_routed_graph(out.ctx, storage.gate_hot, storage.up_hot, storage.down_hot, storage.gate_up_hot, @@ -1009,11 +1421,27 @@ static int mmq_safe_sub_batch() { static const int v = [](){ const char * e = std::getenv("DFLASH_MMQ_SUB_BATCH"); if (e) return std::max(1, std::atoi(e)); +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + return 4; +#else return (query_gpu_compute_sm() >= 80) ? 8 : 1; +#endif }(); return v; } +static bool eval_moe_hybrid_remote_cold_batched( + const MoeHybridConfig & cfg, + 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, + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layer); + static bool eval_moe_hybrid_ffn_batched_core( ggml_backend_t gpu_backend, ggml_backend_t cpu_backend, @@ -1027,7 +1455,9 @@ static bool eval_moe_hybrid_ffn_batched_core( std::vector & out, std::string * err, ggml_gallocr_t * p_hot_alloc, - ggml_gallocr_t * p_cold_alloc) { + ggml_gallocr_t * p_cold_alloc, + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layer) { const int n_embd = cfg.n_embd; const int n_used = cfg.n_expert_used; @@ -1043,36 +1473,34 @@ static bool eval_moe_hybrid_ffn_batched_core( // path below. if (n_tokens > 0 && n_tokens < MoeHybridLayerStorage::kMaxBatchedCache) { const int total_slots = n_used * n_tokens; - const int n_hot_stack = storage.gate_up_hot ? (int)storage.gate_up_hot->ne[2] - : storage.gate_hot ? (int)storage.gate_hot->ne[2] : 1; - const int n_cold_stack = std::max(1, (int)(storage.down_cold ? storage.down_cold->ne[2] : 1)); - std::vector hot_sel(total_slots); - std::vector hot_wts(total_slots, 0.0f); - std::vector cold_sel(total_slots); - std::vector cold_wts(total_slots, 0.0f); - // Dummy (unused) slots must point at INITIALIZED pinned experts, not - // uninitialized cache-ring spare slots (hot_active..n_hot_stack), whose - // garbage Q4_K scale bits could dequantize to NaN (x weight 0 = NaN). - const int n_hot_init = std::max(1, storage.hot_active); - for (int i = 0; i < total_slots; ++i) { hot_sel[i] = i % n_hot_init; cold_sel[i] = i % n_cold_stack; } + auto & scratch = g_mixed_routing_scratch; + scratch.ensure(total_slots, cfg.n_expert); + auto & hot_sel = scratch.hot_sel; + auto & hot_wts = scratch.hot_wts; + auto & cold_sel = scratch.cold_sel; + auto & cold_wts = scratch.cold_wts; bool fp_has_cold = false; - for (int i = 0; i < total_slots; ++i) { - const int32_t gid = selected_ids[i]; - if (gid < 0 || gid >= (int32_t)storage.hot_local_by_global.size()) continue; - const int32_t hl = storage.hot_local_by_global[(size_t)gid]; - if (hl >= 0) { hot_sel[i] = hl; hot_wts[i] = selected_weights[i]; } - else { - const int32_t cl = storage.cold_local_by_global[(size_t)gid]; - if (cl >= 0) { cold_sel[i] = cl; cold_wts[i] = selected_weights[i]; fp_has_cold = true; } - } + bool fp_has_hot = false; + if (!prepare_unique_hot_local_slots( + storage, selected_ids, selected_weights, + n_tokens, n_used, hot_sel, hot_wts, + &fp_has_hot, &fp_has_cold, err)) { + return false; + } + if (fp_has_cold && !prepare_unique_cold_local_slots( + storage, selected_ids, selected_weights, + n_tokens, n_used, cold_sel, cold_wts, + nullptr, err)) { + return false; } CachedHotBatchedGraph & hg = storage.hot_batched_mixed[n_tokens]; - const bool hg_ok = (hg.valid() && hg.n_tokens == n_tokens) + const bool hg_ok = (hg.valid() && hg.n_tokens == n_tokens && !hg.gpu_remap) || build_cached_hot_batched_graph(hg, gpu_backend, storage, desc, cfg, n_tokens); + const bool use_remote_cold = fp_has_cold && expert_compute && expert_layer; CachedHotBatchedGraph * cg = nullptr; bool cg_ok = true; - if (fp_has_cold) { + if (fp_has_cold && !use_remote_cold) { cg = &storage.cold_batched_mixed[n_tokens]; cg_ok = (cg->valid() && cg->n_tokens == n_tokens) || build_cached_cold_batched_graph(*cg, cpu_backend, storage, desc, cfg, n_tokens); @@ -1087,22 +1515,33 @@ static bool eval_moe_hybrid_ffn_batched_core( ggml_backend_graph_compute_async(gpu_backend, hg.gf); std::vector cold_partial; - if (cg) { + if (fp_has_cold) { cold_partial.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); - ggml_backend_tensor_set(cg->inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); - ggml_backend_tensor_set(cg->sel, cold_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); - ggml_backend_tensor_set(cg->wts, cold_wts.data(), 0, sizeof(float) * (size_t)total_slots); - if (ggml_backend_graph_compute(cpu_backend, cg->gf) != GGML_STATUS_SUCCESS) { - ggml_backend_synchronize(gpu_backend); - if (err) *err = "batched cold cached compute failed"; - return false; + if (use_remote_cold) { + if (!eval_moe_hybrid_remote_cold_batched( + cfg, storage, cur_host, selected_ids, + selected_weights, n_tokens, cold_partial, err, + expert_compute, expert_layer)) { + ggml_backend_synchronize(gpu_backend); + if (err && err->empty()) *err = "batched remote cold compute failed"; + return false; + } + } else { + ggml_backend_tensor_set(cg->inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + ggml_backend_tensor_set(cg->sel, cold_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); + ggml_backend_tensor_set(cg->wts, cold_wts.data(), 0, sizeof(float) * (size_t)total_slots); + if (ggml_backend_graph_compute(cpu_backend, cg->gf) != GGML_STATUS_SUCCESS) { + ggml_backend_synchronize(gpu_backend); + if (err) *err = "batched cold cached compute failed"; + return false; + } + ggml_backend_tensor_get(cg->output, cold_partial.data(), 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); } - ggml_backend_tensor_get(cg->output, cold_partial.data(), 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); } ggml_backend_synchronize(gpu_backend); ggml_backend_tensor_get(hg.output, out.data(), 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); - if (cg) { + if (fp_has_cold) { const size_t ntot = (size_t)n_embd * (size_t)n_tokens; for (size_t i = 0; i < ntot; ++i) out[i] += cold_partial[i]; } @@ -1115,35 +1554,24 @@ static bool eval_moe_hybrid_ffn_batched_core( // Dummy slots use weight 0.0 and are distributed evenly across all experts // to avoid pathological routing imbalance that triggers OOB in MMQ stream-k. const int total_slots = n_used * n_tokens; - const int n_hot_stack = storage.gate_up_hot ? (int)storage.gate_up_hot->ne[2] - : storage.gate_hot ? (int)storage.gate_hot->ne[2] - : 1; - std::vector hot_sel(total_slots); - // Dummy slots -> pinned (initialized) experts only; see note above. - const int n_hot_init = std::max(1, storage.hot_active); - for (int i = 0; i < total_slots; ++i) hot_sel[i] = i % n_hot_init; - std::vector hot_wts(total_slots, 0.0f); - std::vector cold_sel(total_slots); - for (int i = 0; i < total_slots; ++i) cold_sel[i] = i % std::max(1, (int)(storage.down_cold ? storage.down_cold->ne[2] : 1)); - std::vector cold_wts(total_slots, 0.0f); + auto & scratch = g_mixed_routing_scratch; + scratch.ensure(total_slots, cfg.n_expert); + auto & hot_sel = scratch.hot_sel; + auto & hot_wts = scratch.hot_wts; + auto & cold_sel = scratch.cold_sel; + auto & cold_wts = scratch.cold_wts; bool has_hot = false, has_cold = false; - - for (int i = 0; i < total_slots; ++i) { - const int32_t gid = selected_ids[i]; - if (gid < 0 || gid >= (int32_t)storage.hot_local_by_global.size()) continue; - const int32_t hot_lid = storage.hot_local_by_global[(size_t)gid]; - if (hot_lid >= 0) { - hot_sel[i] = hot_lid; - hot_wts[i] = selected_weights[i]; - has_hot = true; - } else { - const int32_t cold_lid = storage.cold_local_by_global[(size_t)gid]; - if (cold_lid >= 0) { - cold_sel[i] = cold_lid; - cold_wts[i] = selected_weights[i]; - has_cold = true; - } - } + if (!prepare_unique_hot_local_slots( + storage, selected_ids, selected_weights, + n_tokens, n_used, hot_sel, hot_wts, + &has_hot, &has_cold, err)) { + return false; + } + if (has_cold && !prepare_unique_cold_local_slots( + storage, selected_ids, selected_weights, + n_tokens, n_used, cold_sel, cold_wts, + nullptr, err)) { + return false; } // ── Step 2: Build and run hot GPU graph (includes shared expert always) ── @@ -1219,76 +1647,89 @@ static bool eval_moe_hybrid_ffn_batched_core( } // ── Step 3: Build and run cold CPU graph (overlaps with GPU) ── - std::vector cold_partial((size_t)n_embd * (size_t)n_tokens, 0.0f); + auto & split_scratch = g_hybrid_split_scratch; + auto & cold_partial = split_scratch.cold_partial; + cold_partial.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); if (has_cold) { - ggml_init_params ip{}; - ip.mem_size = 128 * 1024 * 1024; - ip.mem_buffer = nullptr; - ip.no_alloc = true; - ggml_context * cold_ctx = ggml_init(ip); - if (!cold_ctx) { - if (hot_async_launched) ggml_backend_synchronize(gpu_backend); - if (hot_alloc) ggml_gallocr_free(hot_alloc); - if (hot_ctx) ggml_free(hot_ctx); - if (err) *err = "cold ggml_init failed"; - return false; - } - - ggml_tensor * inp = ggml_new_tensor_2d(cold_ctx, GGML_TYPE_F32, n_embd, n_tokens); - ggml_set_input(inp); - ggml_tensor * sel = ggml_new_tensor_2d(cold_ctx, GGML_TYPE_I32, n_used, n_tokens); - ggml_set_input(sel); - ggml_tensor * wts = ggml_new_tensor_2d(cold_ctx, GGML_TYPE_F32, n_used, n_tokens); - ggml_set_input(wts); - - ggml_tensor * cold_routed = nullptr; - build_batched_routed_graph(cold_ctx, - storage.gate_cold, storage.up_cold, storage.down_cold, storage.gate_up_cold, - desc.ffn_gate_exps_s, desc.ffn_up_exps_s, desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, - inp, sel, wts, n_embd, n_ff_exp, n_used, n_tokens, &cold_routed); - - ggml_cgraph * cold_gf = ggml_new_graph_custom(cold_ctx, 4096, false); - ggml_set_output(cold_routed); - ggml_build_forward_expand(cold_gf, cold_routed); - ggml_gallocr_t cold_alloc; - if (p_cold_alloc) { - if (!*p_cold_alloc) - *p_cold_alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(cpu_backend)); - cold_alloc = *p_cold_alloc; + if (expert_compute && expert_layer) { + if (!eval_moe_hybrid_remote_cold_batched( + cfg, storage, cur_host, selected_ids, selected_weights, + n_tokens, cold_partial, err, expert_compute, expert_layer)) { + if (hot_async_launched) ggml_backend_synchronize(gpu_backend); + if (!p_hot_alloc && hot_alloc) ggml_gallocr_free(hot_alloc); + if (hot_ctx) ggml_free(hot_ctx); + if (err && err->empty()) *err = "hybrid batched remote cold compute failed"; + return false; + } } else { - cold_alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(cpu_backend)); - } - if (!ggml_gallocr_alloc_graph(cold_alloc, cold_gf)) { - if (hot_async_launched) ggml_backend_synchronize(gpu_backend); - if (!p_hot_alloc && hot_alloc) ggml_gallocr_free(hot_alloc); - if (hot_ctx) ggml_free(hot_ctx); - if (!p_cold_alloc) ggml_gallocr_free(cold_alloc); - ggml_free(cold_ctx); - if (err) *err = "hybrid batched cold gallocr failed"; - return false; - } + ggml_init_params ip{}; + ip.mem_size = 128 * 1024 * 1024; + ip.mem_buffer = nullptr; + ip.no_alloc = true; + ggml_context * cold_ctx = ggml_init(ip); + if (!cold_ctx) { + if (hot_async_launched) ggml_backend_synchronize(gpu_backend); + if (!p_hot_alloc && hot_alloc) ggml_gallocr_free(hot_alloc); + if (hot_ctx) ggml_free(hot_ctx); + if (err) *err = "cold ggml_init failed"; + return false; + } - ggml_backend_tensor_set(inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); - ggml_backend_tensor_set(sel, cold_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); - ggml_backend_tensor_set(wts, cold_wts.data(), 0, sizeof(float) * (size_t)total_slots); + ggml_tensor * inp = ggml_new_tensor_2d(cold_ctx, GGML_TYPE_F32, n_embd, n_tokens); + ggml_set_input(inp); + ggml_tensor * sel = ggml_new_tensor_2d(cold_ctx, GGML_TYPE_I32, n_used, n_tokens); + ggml_set_input(sel); + ggml_tensor * wts = ggml_new_tensor_2d(cold_ctx, GGML_TYPE_F32, n_used, n_tokens); + ggml_set_input(wts); - // Run CPU synchronously (overlaps with GPU async) - auto st = ggml_backend_graph_compute(cpu_backend, cold_gf); - if (st != GGML_STATUS_SUCCESS) { - if (hot_async_launched) ggml_backend_synchronize(gpu_backend); - if (!p_hot_alloc && hot_alloc) ggml_gallocr_free(hot_alloc); - if (hot_ctx) ggml_free(hot_ctx); + ggml_tensor * cold_routed = nullptr; + build_batched_routed_graph(cold_ctx, + storage.gate_cold, storage.up_cold, storage.down_cold, storage.gate_up_cold, + desc.ffn_gate_exps_s, desc.ffn_up_exps_s, desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, + inp, sel, wts, n_embd, n_ff_exp, n_used, n_tokens, &cold_routed); + + ggml_cgraph * cold_gf = ggml_new_graph_custom(cold_ctx, 4096, false); + ggml_set_output(cold_routed); + ggml_build_forward_expand(cold_gf, cold_routed); + ggml_gallocr_t cold_alloc; + if (p_cold_alloc) { + if (!*p_cold_alloc) + *p_cold_alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(cpu_backend)); + cold_alloc = *p_cold_alloc; + } else { + cold_alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(cpu_backend)); + } + if (!ggml_gallocr_alloc_graph(cold_alloc, cold_gf)) { + if (hot_async_launched) ggml_backend_synchronize(gpu_backend); + if (!p_hot_alloc && hot_alloc) ggml_gallocr_free(hot_alloc); + if (hot_ctx) ggml_free(hot_ctx); + if (!p_cold_alloc) ggml_gallocr_free(cold_alloc); + ggml_free(cold_ctx); + if (err) *err = "hybrid batched cold gallocr failed"; + return false; + } + + ggml_backend_tensor_set(inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + ggml_backend_tensor_set(sel, cold_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); + ggml_backend_tensor_set(wts, cold_wts.data(), 0, sizeof(float) * (size_t)total_slots); + + auto st = ggml_backend_graph_compute(cpu_backend, cold_gf); + if (st != GGML_STATUS_SUCCESS) { + if (hot_async_launched) ggml_backend_synchronize(gpu_backend); + if (!p_hot_alloc && hot_alloc) ggml_gallocr_free(hot_alloc); + if (hot_ctx) ggml_free(hot_ctx); + if (!p_cold_alloc) ggml_gallocr_free(cold_alloc); + ggml_free(cold_ctx); + if (err) *err = "hybrid batched cold compute failed"; + return false; + } + + ggml_backend_tensor_get(cold_routed, cold_partial.data(), 0, + sizeof(float) * (size_t)n_embd * (size_t)n_tokens); if (!p_cold_alloc) ggml_gallocr_free(cold_alloc); ggml_free(cold_ctx); - if (err) *err = "hybrid batched cold compute failed"; - return false; } - - ggml_backend_tensor_get(cold_routed, cold_partial.data(), 0, - sizeof(float) * (size_t)n_embd * (size_t)n_tokens); - if (!p_cold_alloc) ggml_gallocr_free(cold_alloc); - ggml_free(cold_ctx); } // ── Step 4: Sync GPU and read hot result ── @@ -1309,6 +1750,407 @@ static bool eval_moe_hybrid_ffn_batched_core( return true; } +static bool eval_moe_hybrid_remote_cold_batched( + const MoeHybridConfig & cfg, + 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, + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layer) { + const int n_embd = cfg.n_embd; + const int n_used = cfg.n_expert_used; + const int n_ff_exp = cfg.n_ff_exp; + out.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); + if (n_tokens <= 0) return true; + if (!expert_compute || !expert_layer) return true; + + auto & scratch = g_remote_cold_batch_scratch; + scratch.ensure(n_tokens, n_used, n_embd, n_used); + auto & cold_counts = scratch.cold_counts; + auto & cold_sel = scratch.cold_sel; + auto & cold_wts = scratch.cold_wts; + int max_cold_selected = 0; + + for (int t = 0; t < n_tokens; ++t) { + int count = 0; + for (int i = 0; i < n_used; ++i) { + const size_t src = (size_t)t * (size_t)n_used + (size_t)i; + const int32_t gid = selected_ids[src]; + if (gid < 0 || gid >= (int32_t)storage.hot_local_by_global.size()) { + continue; + } + if (storage.hot_local_by_global[(size_t)gid] >= 0) { + continue; + } + const int32_t cold_lid = storage.cold_local_by_global[(size_t)gid]; + if (cold_lid >= 0) { + const size_t dst = (size_t)t * (size_t)n_used + (size_t)count; + cold_sel[dst] = cold_lid; + cold_wts[dst] = selected_weights[src]; + count++; + } + } + cold_counts[(size_t)t] = count; + max_cold_selected = std::max(max_cold_selected, count); + } + + if (max_cold_selected == 0) return true; + if (expert_layer->cold_global_by_local.empty()) { + if (err) *err = "hybrid batched remote cold layer has no cold experts"; + return false; + } + auto & group_ids = scratch.group_ids; + auto & group_wts = scratch.group_wts; + if (group_ids.capacity() < (size_t)n_tokens * (size_t)max_cold_selected) { + group_ids.reserve((size_t)n_tokens * (size_t)max_cold_selected); + } + if (group_wts.capacity() < (size_t)n_tokens * (size_t)max_cold_selected) { + group_wts.reserve((size_t)n_tokens * (size_t)max_cold_selected); + } + + const char * pack_env = std::getenv("DFLASH_MOE_REMOTE_COLD_PACK"); + const bool pack_env_set = pack_env && *pack_env; + const bool padded_pack = pack_env_set + ? (std::strcmp(pack_env, "padded") == 0 || + std::strcmp(pack_env, "single") == 0) + : expert_compute->prefers_padded_batch(); + if (padded_pack) { + const int cold_stack = (int)expert_layer->cold_global_by_local.size(); + if (cold_stack < max_cold_selected) { + if (err) *err = "hybrid batched remote cold padded pack exceeds cold stack"; + return false; + } + group_ids.resize((size_t)n_tokens * (size_t)max_cold_selected); + group_wts.resize((size_t)n_tokens * (size_t)max_cold_selected); + std::vector used((size_t)cold_stack, 0); + for (int t = 0; t < n_tokens; ++t) { + const int count = cold_counts[(size_t)t]; + std::fill(used.begin(), used.end(), 0); + int next_dummy = 0; + for (int i = 0; i < max_cold_selected; ++i) { + const size_t dst = (size_t)t * (size_t)max_cold_selected + (size_t)i; + if (i < count) { + const size_t src = (size_t)t * (size_t)n_used + (size_t)i; + const int32_t cold_local = cold_sel[src]; + if (cold_local < 0 || cold_local >= cold_stack) { + if (err) *err = "hybrid batched remote cold local id out of range"; + return false; + } + if (used[(size_t)cold_local]) { + if (err) *err = "hybrid batched remote cold duplicate local id"; + return false; + } + group_ids[dst] = cold_local; + group_wts[dst] = cold_wts[src]; + used[(size_t)cold_local] = 1; + } else { + while (next_dummy < cold_stack && used[(size_t)next_dummy]) { + ++next_dummy; + } + if (next_dummy >= cold_stack) { + if (err) *err = "hybrid batched remote cold lacks unique padded ids"; + return false; + } + group_ids[dst] = next_dummy; + group_wts[dst] = 0.0f; + used[(size_t)next_dummy] = 1; + ++next_dummy; + } + } + } + if (!expert_compute->compute_batch(*expert_layer, + cur_host, + group_ids.data(), + group_wts.data(), + n_tokens, + max_cold_selected, + n_embd, n_ff_exp, + out.data())) { + if (err) *err = "hybrid batched remote cold padded compute failed"; + return false; + } + return true; + } + + auto & token_groups = scratch.token_groups; + if (token_groups.size() < (size_t)max_cold_selected + 1) { + token_groups.resize((size_t)max_cold_selected + 1); + } + auto & active_groups = scratch.active_groups; + for (int n_cold : active_groups) { + if (n_cold >= 0 && (size_t)n_cold < token_groups.size()) { + token_groups[(size_t)n_cold].clear(); + } + } + active_groups.clear(); + for (int t = 0; t < n_tokens; ++t) { + const int count = cold_counts[(size_t)t]; + if (count > 0) { + auto & bucket = token_groups[(size_t)count]; + if (bucket.empty()) { + active_groups.push_back(count); + } + bucket.push_back(t); + } + } + + auto & group_input = scratch.group_input; + auto & group_output = scratch.group_output; + if (group_input.capacity() < (size_t)n_embd * (size_t)n_tokens) { + group_input.reserve((size_t)n_embd * (size_t)n_tokens); + } + if (group_output.capacity() < (size_t)n_embd * (size_t)n_tokens) { + group_output.reserve((size_t)n_embd * (size_t)n_tokens); + } + for (int n_cold : active_groups) { + const auto & token_group = token_groups[(size_t)n_cold]; + if (token_group.empty()) continue; + const size_t group_limit = (size_t)std::max( + 1, std::min( + moe_expert_compute_batch_limit_from_env(), + (int)token_group.size())); + for (size_t base = 0; base < token_group.size(); base += group_limit) { + const int tc = (int)std::min(group_limit, token_group.size() - base); + bool contiguous_tokens = true; + for (int gi = 1; gi < tc; ++gi) { + if (token_group[base + (size_t)gi] != token_group[base] + gi) { + contiguous_tokens = false; + break; + } + } + + const int first_token = token_group[base]; + const float * compute_input = contiguous_tokens + ? cur_host + (size_t)first_token * (size_t)n_embd + : nullptr; + float * compute_output = contiguous_tokens + ? out.data() + (size_t)first_token * (size_t)n_embd + : nullptr; + if (!contiguous_tokens) { + group_input.resize((size_t)tc * (size_t)n_embd); + group_output.resize((size_t)tc * (size_t)n_embd); + compute_input = group_input.data(); + compute_output = group_output.data(); + } + + group_ids.resize((size_t)tc * (size_t)n_cold); + group_wts.resize((size_t)tc * (size_t)n_cold); + for (int gi = 0; gi < tc; ++gi) { + const int t = token_group[base + (size_t)gi]; + if (!contiguous_tokens) { + std::memcpy(group_input.data() + (size_t)gi * (size_t)n_embd, + cur_host + (size_t)t * (size_t)n_embd, + sizeof(float) * (size_t)n_embd); + } + for (int i = 0; i < n_cold; ++i) { + const size_t src = (size_t)t * (size_t)n_used + (size_t)i; + const size_t dst = (size_t)gi * (size_t)n_cold + (size_t)i; + group_ids[dst] = cold_sel[src]; + group_wts[dst] = cold_wts[src]; + } + } + + if (!expert_compute->compute_batch(*expert_layer, + compute_input, + group_ids.data(), + group_wts.data(), + tc, n_cold, + n_embd, n_ff_exp, + compute_output)) { + if (err) *err = "hybrid batched remote cold compute failed"; + return false; + } + + if (!contiguous_tokens) { + for (int gi = 0; gi < tc; ++gi) { + const int t = token_group[base + (size_t)gi]; + std::memcpy(out.data() + (size_t)t * (size_t)n_embd, + group_output.data() + (size_t)gi * (size_t)n_embd, + sizeof(float) * (size_t)n_embd); + } + } + } + token_groups[(size_t)n_cold].clear(); + } + return true; +} + +static bool eval_moe_hot_remap_batched( + ggml_backend_t gpu_backend, + const MoeHybridConfig & cfg, + const MoeLayerDesc & desc, + MoeHybridLayerStorage & storage, + const float * cur_host, + const int32_t * selected_ids, + const float * selected_weights, + int n_tokens, + std::vector & out, + std::string * err) { + const int n_embd = cfg.n_embd; + 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 (n_tokens <= 0) return true; + if (!cur_host || !selected_ids || !selected_weights) return false; + if ((int)storage.hot_local_by_global.size() < cfg.n_expert) { + if (err) *err = "hot-remap batched graph lacks hot-local LUT"; + return false; + } + + auto & cached = (n_tokens > 0 && n_tokens < MoeHybridLayerStorage::kMaxBatchedCache) + ? storage.hot_batched_mixed[n_tokens] + : storage.hot_batched_graph; + if (!cached.valid() || cached.n_tokens != n_tokens || !cached.gpu_remap || + !cached.global_ids || !cached.raw_weights || + !cached.hot_local_lut || !cached.valid_lut) { + if (!build_cached_hot_batched_graph(cached, gpu_backend, storage, + desc, cfg, n_tokens, + /*gpu_remap=*/true)) { + if (err) *err = "hot-remap batched graph build failed"; + return false; + } + } + + auto & scratch = g_mixed_routing_scratch; + scratch.ensure(total_slots, cfg.n_expert); + auto & lut = scratch.lut; + auto & vlut = scratch.vlut; + for (int expert = 0; expert < cfg.n_expert; ++expert) { + const int32_t local = storage.hot_local_by_global[(size_t)expert]; + if (local >= 0) { + lut[(size_t)expert] = local; + vlut[(size_t)expert] = 1.0f; + } else { + lut[(size_t)expert] = 0; + vlut[(size_t)expert] = 0.0f; + } + } + + ggml_backend_tensor_set(cached.inp, cur_host, 0, + sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + ggml_backend_tensor_set(cached.hot_local_lut, lut.data(), 0, + sizeof(int32_t) * (size_t)cfg.n_expert); + ggml_backend_tensor_set(cached.valid_lut, vlut.data(), 0, + sizeof(float) * (size_t)cfg.n_expert); + ggml_backend_tensor_set(cached.global_ids, selected_ids, 0, + sizeof(int32_t) * (size_t)total_slots); + ggml_backend_tensor_set(cached.raw_weights, selected_weights, 0, + sizeof(float) * (size_t)total_slots); + + if (ggml_backend_graph_compute(gpu_backend, cached.gf) != GGML_STATUS_SUCCESS) { + if (err) *err = "hot-remap batched graph compute failed"; + return false; + } + ggml_backend_tensor_get(cached.output, out.data(), 0, + sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + return true; +} + +static bool eval_moe_hybrid_ffn_split_hot_remote_cold( + ggml_backend_t gpu_backend, + const MoeHybridConfig & cfg, + const MoeLayerDesc & desc, + MoeHybridLayerStorage & storage, + const float * cur_host, + const int32_t * selected_ids, + const float * selected_weights, + int n_tokens, + int hot_sub_batch, + std::vector & out, + std::string * err, + ggml_gallocr_t * p_hot_alloc, + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layer) { + const int n_embd = cfg.n_embd; + const int n_used = cfg.n_expert_used; + std::fprintf(stderr, + "[hybrid-ffn] split_hot_remote_cold n_used=%d n_tokens=%d hot_sub_batch=%d\n", + n_used, n_tokens, hot_sub_batch); + out.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); + const bool profile_enabled = moe_hybrid_batched_profile_enabled(); + const auto total_t0 = profile_enabled ? HybridClock::now() : HybridClock::time_point{}; + + auto & split_scratch = g_hybrid_split_scratch; + auto & cold_partial = split_scratch.cold_partial; + cold_partial.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); + const auto cold_t0 = profile_enabled ? HybridClock::now() : HybridClock::time_point{}; + bool cold_ok = false; + std::string cold_err; + g_hybrid_cold_worker.submit([&]() { + cold_ok = eval_moe_hybrid_remote_cold_batched( + cfg, storage, cur_host, selected_ids, selected_weights, + n_tokens, cold_partial, &cold_err, expert_compute, expert_layer); + }); + + auto & hot_out = split_scratch.hot_out; + bool hot_ok = true; + std::string hot_err; + uint64_t local_hot_calls = 0; + uint64_t local_hot_tokens = 0; + uint64_t local_hot_us = 0; + for (int t0 = 0; t0 < n_tokens; t0 += hot_sub_batch) { + const int tc = std::min(hot_sub_batch, n_tokens - t0); + const auto hot_t0 = profile_enabled ? HybridClock::now() : HybridClock::time_point{}; + if (!eval_moe_hot_remap_batched( + gpu_backend, cfg, desc, storage, + cur_host + (size_t)t0 * (size_t)n_embd, + selected_ids + (size_t)t0 * (size_t)n_used, + selected_weights + (size_t)t0 * (size_t)n_used, + tc, hot_out, &hot_err)) { + hot_ok = false; + break; + } + const auto hot_t1 = profile_enabled ? HybridClock::now() : HybridClock::time_point{}; + if (profile_enabled) { + local_hot_calls++; + local_hot_tokens += (uint64_t)tc; + local_hot_us += elapsed_us(hot_t0, hot_t1); + } + std::memcpy(out.data() + (size_t)t0 * (size_t)n_embd, + hot_out.data(), + sizeof(float) * (size_t)n_embd * (size_t)tc); + } + + const auto overlap_wait_t0 = profile_enabled ? HybridClock::now() : HybridClock::time_point{}; + g_hybrid_cold_worker.wait(); + const auto overlap_wait_t1 = profile_enabled ? HybridClock::now() : HybridClock::time_point{}; + const auto cold_t1 = profile_enabled ? HybridClock::now() : HybridClock::time_point{}; + if (!hot_ok) { + if (err) *err = hot_err.empty() ? "hybrid batched hot compute failed" : hot_err; + return false; + } + if (!cold_ok) { + if (err) *err = cold_err.empty() ? "hybrid batched remote cold compute failed" : cold_err; + return false; + } + + const auto merge_t0 = profile_enabled ? HybridClock::now() : HybridClock::time_point{}; + const size_t total_floats = (size_t)n_embd * (size_t)n_tokens; + for (size_t i = 0; i < total_floats; ++i) { + out[i] += cold_partial[i]; + } + const auto merge_t1 = profile_enabled ? HybridClock::now() : HybridClock::time_point{}; + if (profile_enabled) { + const auto total_t1 = HybridClock::now(); + auto & prof = moe_hybrid_batched_profile(); + prof.calls++; + prof.tokens += (uint64_t)n_tokens; + prof.remote_cold_us += elapsed_us(cold_t0, cold_t1); + prof.hot_calls += local_hot_calls; + prof.hot_tokens += local_hot_tokens; + prof.hot_us += local_hot_us; + prof.overlap_us += elapsed_us(overlap_wait_t0, overlap_wait_t1); + prof.merge_us += elapsed_us(merge_t0, merge_t1); + prof.total_us += elapsed_us(total_t0, total_t1); + prof.maybe_print(false); + } + return true; +} + // ── Hot-Only Batched Prefill ── // When all selected experts are in VRAM, skip cold entirely: no CPU graph, // no partition into hot/cold, no merge loop. Pure GPU. @@ -1332,6 +2174,18 @@ bool eval_moe_hot_only_batched( out.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); if (n_tokens <= 0) return true; + // The batched hot-only fast path has only been validated on the legacy + // top-k<=8 shapes. Larger routed top-k values, such as Ornith's + // n_expert_used=10, should stay on the hybrid path until their mul_mat_id + // shapes are explicitly verified. + if (n_used > 8) { + std::fprintf(stderr, + "[hybrid-ffn] reject hot_only batched path n_used=%d n_tokens=%d\n", + n_used, n_tokens); + if (err) *err = "hot_only batched fast path disabled for n_expert_used > 8"; + return false; + } + // Workaround for the ggml-cuda MMQ mul_mat_id stream-k fault on a REDUCED // hot stack (sm_75/gfx1151 AND sm_86): slice sub-64 batches to a size the // MMVQ-mmid path handles. See mmq_safe_sub_batch(). @@ -1359,18 +2213,6 @@ bool eval_moe_hot_only_batched( return true; } - // Remap global expert IDs → hot-local IDs - const int total_slots = n_used * n_tokens; - std::vector hot_sel(total_slots); - for (int i = 0; i < total_slots; ++i) { - const int32_t gid = selected_ids[i]; - if (gid < 0 || gid >= (int32_t)storage.hot_local_by_global.size()) { - hot_sel[i] = 0; - } else { - hot_sel[i] = storage.hot_local_by_global[(size_t)gid]; - } - } - // ── Fast path: use cached graph (avoids rebuild + realloc) ── // Per-n_tokens cache: spec-decode alternates verify (verify_width) and // replay (commit_n) sizes; a single slot would rebuild all 40 layers' FFN @@ -1378,11 +2220,67 @@ bool eval_moe_hot_only_batched( auto & cached = (n_tokens > 0 && n_tokens < MoeHybridLayerStorage::kMaxBatchedCache) ? storage.hot_batched_mixed[n_tokens] : storage.hot_batched_graph; - if (cached.n_tokens == n_tokens && cached.valid()) { + const bool use_remap = cached.valid() && cached.gpu_remap; + if (cached.n_tokens == n_tokens && cached.valid() && use_remap) { + const int total_slots = n_used * n_tokens; + auto & scratch = g_mixed_routing_scratch; + scratch.ensure(total_slots, cfg.n_expert); + auto & lut = scratch.lut; + auto & vlut = scratch.vlut; + auto & hot_sel = scratch.hot_sel; + auto & hot_wts = scratch.hot_wts; + for (int expert = 0; expert < cfg.n_expert; ++expert) { + const int32_t local = storage.hot_local_by_global[(size_t)expert]; + if (local >= 0) { + lut[(size_t)expert] = local; + vlut[(size_t)expert] = 1.0f; + } else { + lut[(size_t)expert] = 0; + vlut[(size_t)expert] = 0.0f; + } + } + ggml_backend_tensor_set(cached.inp, cur_host, 0, + sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + ggml_backend_tensor_set(cached.hot_local_lut, lut.data(), 0, + sizeof(int32_t) * (size_t)cfg.n_expert); + ggml_backend_tensor_set(cached.valid_lut, vlut.data(), 0, + sizeof(float) * (size_t)cfg.n_expert); + if (!prepare_unique_hot_global_slots( + storage, selected_ids, selected_weights, + n_tokens, n_used, hot_sel, hot_wts, err)) { + return false; + } + ggml_backend_tensor_set(cached.global_ids, hot_sel.data(), 0, + sizeof(int32_t) * (size_t)total_slots); + ggml_backend_tensor_set(cached.raw_weights, hot_wts.data(), 0, + sizeof(float) * (size_t)total_slots); + auto st = ggml_backend_graph_compute(gpu_backend, cached.gf); + if (st != GGML_STATUS_SUCCESS) { + if (err) *err = "hot_only remap cached compute failed"; + return false; + } + ggml_backend_tensor_get(cached.output, out.data(), 0, + sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + return true; + } + if (cached.n_tokens == n_tokens && cached.valid() && !cached.gpu_remap) { // Reuse pre-built graph: just upload data and compute + const int total_slots = n_used * n_tokens; + auto & scratch = g_mixed_routing_scratch; + scratch.ensure(total_slots, cfg.n_expert); + auto & hot_sel = scratch.hot_sel; + auto & hot_wts = scratch.hot_wts; + bool has_hot = false; + bool has_cold = false; + if (!prepare_unique_hot_local_slots( + storage, selected_ids, selected_weights, + n_tokens, n_used, hot_sel, hot_wts, + &has_hot, &has_cold, err)) { + return false; + } ggml_backend_tensor_set(cached.inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); ggml_backend_tensor_set(cached.sel, hot_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); - ggml_backend_tensor_set(cached.wts, selected_weights, 0, sizeof(float) * (size_t)total_slots); + ggml_backend_tensor_set(cached.wts, hot_wts.data(), 0, sizeof(float) * (size_t)total_slots); auto st = ggml_backend_graph_compute(gpu_backend, cached.gf); if (st != GGML_STATUS_SUCCESS) { @@ -1398,19 +2296,13 @@ bool eval_moe_hot_only_batched( // Cache when: sub-batch size (legacy), full stack (all hot), or full-batch safe (sm_80+). if (mmq_full_batch_ok(cfg, n_tokens) || n_tokens <= MMQ_SAFE_SUB_BATCH || (n_hot_stack == 0 || n_hot_stack >= cfg.n_expert)) { - if (build_cached_hot_batched_graph(cached, gpu_backend, storage, desc, cfg, n_tokens)) { + if (build_cached_hot_batched_graph(cached, gpu_backend, storage, desc, cfg, n_tokens, + /*gpu_remap=*/true)) { // Successfully cached — use it immediately - ggml_backend_tensor_set(cached.inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); - ggml_backend_tensor_set(cached.sel, hot_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); - ggml_backend_tensor_set(cached.wts, selected_weights, 0, sizeof(float) * (size_t)total_slots); - - auto st = ggml_backend_graph_compute(gpu_backend, cached.gf); - if (st != GGML_STATUS_SUCCESS) { - if (err) *err = "hot_only cached compute failed (first)"; - return false; - } - ggml_backend_tensor_get(cached.output, out.data(), 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); - return true; + return eval_moe_hot_remap_batched(gpu_backend, cfg, desc, storage, + cur_host, selected_ids, + selected_weights, n_tokens, + out, err); } // Fall through to uncached path if build fails } @@ -1469,8 +2361,30 @@ bool eval_moe_hot_only_batched( } ggml_backend_tensor_set(inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); - ggml_backend_tensor_set(sel, hot_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); - ggml_backend_tensor_set(wts, selected_weights, 0, sizeof(float) * (size_t)total_slots); + if (!cached.gpu_remap) { + const int total_slots = n_used * n_tokens; + auto & scratch = g_mixed_routing_scratch; + scratch.ensure(total_slots, cfg.n_expert); + auto & hot_sel = scratch.hot_sel; + auto & hot_wts = scratch.hot_wts; + bool has_hot = false; + bool has_cold = false; + if (!prepare_unique_hot_local_slots( + storage, selected_ids, selected_weights, + n_tokens, n_used, hot_sel, hot_wts, + &has_hot, &has_cold, err)) { + if (!p_hot_alloc) ggml_gallocr_free(alloc); + ggml_free(ctx); + return false; + } + ggml_backend_tensor_set(sel, hot_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); + ggml_backend_tensor_set(wts, hot_wts.data(), 0, sizeof(float) * (size_t)total_slots); + } else { + if (err) *err = "hot_only remap cached graph should be handled before uncached fallback"; + if (!p_hot_alloc) ggml_gallocr_free(alloc); + ggml_free(ctx); + return false; + } auto st = ggml_backend_graph_compute(gpu_backend, gf); if (st != GGML_STATUS_SUCCESS) { @@ -1508,13 +2422,33 @@ bool eval_moe_hybrid_ffn_batched( std::vector & out, std::string * err, ggml_gallocr_t * p_hot_alloc, - ggml_gallocr_t * p_cold_alloc) { + ggml_gallocr_t * p_cold_alloc, + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layer) { + if (cfg.n_expert_used > 8) { + std::fprintf(stderr, + "[hybrid-ffn] disable batched path for n_used=%d n_tokens=%d\n", + cfg.n_expert_used, n_tokens); + if (err) *err = "hybrid batched path disabled for n_expert_used > 8"; + return false; + } const int n_hot_stack = storage.gate_up_hot ? (int)storage.gate_up_hot->ne[2] : storage.gate_hot ? (int)storage.gate_hot->ne[2] : 0; const int MMQ_SAFE_SUB_BATCH = mmq_safe_sub_batch(); + std::fprintf(stderr, + "[hybrid-ffn] batched entry n_used=%d n_tokens=%d n_hot_stack=%d mmq_safe=%d remote_cold=%d\n", + cfg.n_expert_used, n_tokens, n_hot_stack, MMQ_SAFE_SUB_BATCH, + (expert_compute && expert_layer) ? 1 : 0); if (!mmq_full_batch_ok(cfg, n_tokens) && n_hot_stack > 0 && n_tokens > MMQ_SAFE_SUB_BATCH) { + if (expert_compute && expert_layer) { + return eval_moe_hybrid_ffn_split_hot_remote_cold( + gpu_backend, cfg, desc, storage, + cur_host, selected_ids, selected_weights, + n_tokens, MMQ_SAFE_SUB_BATCH, out, err, p_hot_alloc, + expert_compute, expert_layer); + } const int n_embd = cfg.n_embd; const int n_used = cfg.n_expert_used; out.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); @@ -1526,7 +2460,8 @@ bool eval_moe_hybrid_ffn_batched( cur_host + (size_t)t0 * (size_t)n_embd, selected_ids + (size_t)t0 * (size_t)n_used, selected_weights + (size_t)t0 * (size_t)n_used, - tc, sub_out, err, p_hot_alloc, p_cold_alloc)) { + tc, sub_out, err, p_hot_alloc, p_cold_alloc, + expert_compute, expert_layer)) { return false; } std::memcpy(out.data() + (size_t)t0 * (size_t)n_embd, @@ -1538,7 +2473,7 @@ bool eval_moe_hybrid_ffn_batched( return eval_moe_hybrid_ffn_batched_core( gpu_backend, cpu_backend, cfg, desc, storage, cur_host, selected_ids, selected_weights, n_tokens, out, err, - p_hot_alloc, p_cold_alloc); + p_hot_alloc, p_cold_alloc, expert_compute, expert_layer); } void ResidualCombineGraph::free() { @@ -1710,8 +2645,10 @@ bool eval_moe_hybrid_ffn_gpu_resident( return false; } { - std::vector lut((size_t)cfg.n_expert); - std::vector vlut((size_t)cfg.n_expert); + auto & scratch = g_mixed_routing_scratch; + scratch.ensure(n_selected, cfg.n_expert); + auto & lut = scratch.lut; + auto & vlut = scratch.vlut; for (int e = 0; e < cfg.n_expert; ++e) { const int32_t l = storage.hot_local_by_global[(size_t)e]; lut[(size_t)e] = (l >= 0) ? l : 0; diff --git a/server/src/common/moe_hybrid_ffn_eval.h b/server/src/common/moe_hybrid_ffn_eval.h index 591017bba..665f46bdb 100644 --- a/server/src/common/moe_hybrid_ffn_eval.h +++ b/server/src/common/moe_hybrid_ffn_eval.h @@ -4,6 +4,7 @@ #include "moe_hybrid_types.h" #include "moe_hybrid_storage.h" +#include "moe_expert_compute.h" #include "ggml-backend.h" @@ -138,7 +139,9 @@ bool eval_moe_hybrid_ffn_batched( std::vector & out, std::string * err = nullptr, ggml_gallocr_t * p_hot_alloc = nullptr, - ggml_gallocr_t * p_cold_alloc = nullptr); + ggml_gallocr_t * p_cold_alloc = nullptr, + MoeExpertCompute * expert_compute = nullptr, + const MoeExpertLayer * expert_layer = nullptr); // Hot-only batched prefill: all selected experts are in VRAM. // Skips cold graph build, CPU compute, and merge — pure GPU path. @@ -212,6 +215,7 @@ bool build_cached_hot_batched_graph( MoeHybridLayerStorage & storage, const MoeLayerDesc & desc, const MoeHybridConfig & cfg, - int n_tokens); + int n_tokens, + bool gpu_remap = false); } // namespace dflash::common diff --git a/server/src/common/moe_hybrid_placement.cpp b/server/src/common/moe_hybrid_placement.cpp index 49bcd5ea1..ad67dc04e 100644 --- a/server/src/common/moe_hybrid_placement.cpp +++ b/server/src/common/moe_hybrid_placement.cpp @@ -1,4 +1,5 @@ #include "moe_hybrid_placement.h" +#include "expert_split_plan.h" #include "moe_hybrid_routing_stats.h" #include @@ -9,6 +10,57 @@ namespace dflash::common { +namespace { + +bool build_plan_from_stats_with_layer_bytes( + const MoeHybridRoutingStats & stats, + const std::vector & layer_expert_bytes, + uint64_t total_hot_budget_bytes, + int min_hot_per_layer, + ExpertSplitPlan & out, + std::string * err) { + if (stats.empty() || stats.n_layer <= 0 || stats.n_expert <= 0) { + if (err) *err = "stats not initialized"; + return false; + } + if ((int) layer_expert_bytes.size() != stats.n_layer) { + if (err) *err = "layer_expert_bytes size mismatch"; + return false; + } + if (total_hot_budget_bytes == 0) { + if (err) *err = "total_hot_budget_bytes must be > 0"; + return false; + } + + ExpertSplitConfig cfg; + cfg.n_layer = stats.n_layer; + cfg.n_expert = stats.n_expert; + cfg.allow_implicit_cpu_fallback = true; + cfg.require_full_grid = true; + cfg.min_per_layer_by_target = {std::max(0, std::min(min_hot_per_layer, stats.n_expert))}; + + std::vector targets = { + {"hot", "gpu", 0, total_hot_budget_bytes, 0, false}, + }; + + std::vector units; + units.reserve((size_t) stats.n_layer * (size_t) stats.n_expert); + for (int il = 0; il < stats.n_layer; ++il) { + for (int ie = 0; ie < stats.n_expert; ++ie) { + units.push_back(ExpertSplitUnit{ + il, + ie, + layer_expert_bytes[(size_t) il], + (double) stats.count(il, ie), + }); + } + } + + return build_expert_split_plan(cfg, targets, units, out, err); +} + +} // namespace + bool MoeHybridPlacement::matches(int n_layer_, int n_expert_, int n_expert_used_) const { return n_layer == n_layer_ && n_expert == n_expert_ && @@ -182,80 +234,46 @@ bool MoeHybridPlacement::build_from_stats_with_layer_bytes( int min_hot_per_layer, MoeHybridPlacement & out, std::string * err) { - if (stats.empty() || stats.n_layer <= 0 || stats.n_expert <= 0) { - if (err) *err = "stats not initialized"; - return false; - } - if ((int)layer_expert_bytes.size() != stats.n_layer) { - if (err) *err = "layer_expert_bytes size mismatch"; - return false; - } if (min_hot_per_layer < 0) min_hot_per_layer = 0; - if (total_hot_budget_bytes == 0) { - if (err) *err = "total_hot_budget_bytes must be > 0"; - return false; - } - const int per_layer_floor = std::min(min_hot_per_layer, stats.n_expert); - uint64_t floor_bytes = 0; - for (int il = 0; il < stats.n_layer; ++il) { - if (layer_expert_bytes[(size_t)il] > 0) - floor_bytes += (uint64_t)per_layer_floor * layer_expert_bytes[(size_t)il]; - } - if (floor_bytes > total_hot_budget_bytes) { - if (err) *err = "min_hot_per_layer exceeds byte budget"; + ExpertSplitPlan plan; + if (!build_plan_from_stats_with_layer_bytes( + stats, layer_expert_bytes, total_hot_budget_bytes, + min_hot_per_layer, plan, err)) { return false; } + const int hot_target_index = 0; MoeHybridPlacement tmp; tmp.n_layer = stats.n_layer; tmp.n_expert = stats.n_expert; tmp.n_expert_used = stats.n_expert_used; - tmp.hot_counts.resize((size_t)tmp.n_layer); - for (int il = 0; il < tmp.n_layer; ++il) { - tmp.hot_counts[(size_t)il] = (layer_expert_bytes[(size_t)il] > 0) ? per_layer_floor : 0; - } + tmp.hot_counts.assign((size_t) tmp.n_layer, 0); + tmp.hot_expert_ids.resize((size_t) tmp.n_layer); - std::vector> ranked((size_t)tmp.n_layer); for (int il = 0; il < tmp.n_layer; ++il) { - ranked[(size_t)il] = stats.ranked_experts(il); - } - - uint64_t remaining = total_hot_budget_bytes - floor_bytes; - while (true) { - int best_layer = -1; - double best_value = -1.0; - uint64_t best_gain = 0; - for (int il = 0; il < tmp.n_layer; ++il) { - const int cur_hot = tmp.hot_counts[(size_t)il]; - if (cur_hot >= tmp.n_expert) continue; - const uint64_t bytes = layer_expert_bytes[(size_t)il]; - if (bytes == 0 || bytes > remaining) continue; - const int next_expert = ranked[(size_t)il][(size_t)cur_hot]; - const uint64_t gain = stats.count(il, next_expert); - const double value = (double)gain / (double)bytes; - if (best_layer < 0 || value > best_value || - (value == best_value && gain > best_gain)) { - best_layer = il; - best_value = value; - best_gain = gain; + std::vector layer_hot; + layer_hot.reserve((size_t) tmp.n_expert); + for (int ie = 0; ie < tmp.n_expert; ++ie) { + if (plan.at(il, ie).target_index == hot_target_index) { + layer_hot.push_back(ie); } } - if (best_layer < 0) break; - tmp.hot_counts[(size_t)best_layer]++; - remaining -= layer_expert_bytes[(size_t)best_layer]; - } - - tmp.total_hot = std::accumulate(tmp.hot_counts.begin(), tmp.hot_counts.end(), 0); - tmp.hot_expert_ids.resize((size_t)tmp.n_layer); - for (int il = 0; il < tmp.n_layer; ++il) { - const int hot_n = tmp.hot_counts[(size_t)il]; - auto & hot = tmp.hot_expert_ids[(size_t)il]; - hot.reserve((size_t)hot_n); - for (int i = 0; i < hot_n; ++i) { - hot.push_back((int32_t)ranked[(size_t)il][(size_t)i]); + std::stable_sort(layer_hot.begin(), layer_hot.end(), + [&](int lhs, int rhs) { + const uint64_t lc = stats.count(il, lhs); + const uint64_t rc = stats.count(il, rhs); + if (lc != rc) return lc > rc; + return lhs < rhs; + }); + tmp.hot_counts[(size_t) il] = (int) layer_hot.size(); + auto & hot = tmp.hot_expert_ids[(size_t) il]; + hot.reserve(layer_hot.size()); + for (int expert : layer_hot) { + hot.push_back((int32_t) expert); } } + tmp.total_hot = std::accumulate(tmp.hot_counts.begin(), tmp.hot_counts.end(), 0); out = std::move(tmp); return true; diff --git a/server/src/common/moe_hybrid_storage.cpp b/server/src/common/moe_hybrid_storage.cpp index ab47752f2..f141e30ad 100644 --- a/server/src/common/moe_hybrid_storage.cpp +++ b/server/src/common/moe_hybrid_storage.cpp @@ -58,7 +58,13 @@ void CachedHotBatchedGraph::free() { sel = nullptr; wts = nullptr; output = nullptr; + global_ids = nullptr; + raw_weights = nullptr; + hot_local_lut = nullptr; + valid_lut = nullptr; + residual_in = nullptr; n_tokens = 0; + gpu_remap = false; } namespace { @@ -113,7 +119,16 @@ static bool validate_expert_tensor(ggml_tensor * tensor, int n_expert, size_t * return true; } if (tensor->ne[2] != n_expert) { - if (err) *err = "tensor expert dimension mismatch"; + if (err) { + char msg[256]; + std::snprintf(msg, sizeof(msg), + "tensor expert dimension mismatch: %s ne=[%lld,%lld,%lld,%lld] expected experts=%d", + tensor->name, + (long long)tensor->ne[0], (long long)tensor->ne[1], + (long long)tensor->ne[2], (long long)tensor->ne[3], + n_expert); + *err = msg; + } return false; } if ((int64_t)tensor->nb[2] <= 0) { @@ -130,6 +145,83 @@ static ggml_tensor * new_like_with_expert_count(ggml_context * ctx, ggml_tensor return ggml_new_tensor(ctx, src->type, 4, ne); } +static bool build_layer_expert_residency( + const MoeHybridConfig & cfg, + const std::vector & hot_expert_ids, + const std::vector * cold_expert_order, + std::vector & hot_local_by_global, + std::vector & cold_local_by_global, + std::vector & cold_expert_ids, + uint64_t expert_vram_mask[4], + std::string * err) { + hot_local_by_global.assign((size_t)cfg.n_expert, -1); + cold_local_by_global.assign((size_t)cfg.n_expert, -1); + cold_expert_ids.clear(); + std::memset(expert_vram_mask, 0, sizeof(uint64_t) * 4); + + std::vector hot_seen((size_t)cfg.n_expert, 0); + for (size_t i = 0; i < hot_expert_ids.size(); ++i) { + const int32_t expert = hot_expert_ids[i]; + if (expert < 0 || expert >= cfg.n_expert) { + if (err) *err = "hot expert id out of range"; + return false; + } + if (hot_seen[(size_t)expert]) { + if (err) *err = "duplicate hot expert id"; + return false; + } + hot_local_by_global[(size_t)expert] = (int32_t)i; + hot_seen[(size_t)expert] = 1; + if (expert < 256) { + expert_vram_mask[expert >> 6] |= (1ULL << (expert & 63)); + } + } + + if (!cold_expert_order) { + for (int expert = 0; expert < cfg.n_expert; ++expert) { + if (!hot_seen[(size_t)expert]) { + cold_local_by_global[(size_t)expert] = + (int32_t)cold_expert_ids.size(); + cold_expert_ids.push_back((int32_t)expert); + } + } + return true; + } + + std::vector cold_seen((size_t)cfg.n_expert, 0); + cold_expert_ids.reserve(cold_expert_order->size()); + for (int32_t expert : *cold_expert_order) { + if (expert < 0 || expert >= cfg.n_expert) { + if (err) *err = "cold expert override id out of range"; + return false; + } + if (hot_seen[(size_t)expert]) { + if (err) *err = "cold expert override includes a hot expert"; + return false; + } + if (cold_seen[(size_t)expert]) { + if (err) *err = "cold expert override contains duplicates"; + return false; + } + cold_seen[(size_t)expert] = 1; + cold_local_by_global[(size_t)expert] = (int32_t)cold_expert_ids.size(); + cold_expert_ids.push_back(expert); + } + + const size_t expected_cold = (size_t)cfg.n_expert - hot_expert_ids.size(); + if (cold_expert_ids.size() != expected_cold) { + if (err) *err = "cold expert override count mismatch"; + return false; + } + for (int expert = 0; expert < cfg.n_expert; ++expert) { + if (!hot_seen[(size_t)expert] && !cold_seen[(size_t)expert]) { + if (err) *err = "cold expert override is missing experts"; + return false; + } + } + return true; +} + } // namespace MoeHybridStorage::~MoeHybridStorage() { @@ -194,7 +286,8 @@ bool build_moe_hybrid_storage(const MoeHybridConfig & cfg, const MoeHybridPlacement & placement, const std::vector & layer_descs, MoeHybridStorage & out, - std::string * err) { + std::string * err, + const std::vector> * cold_expert_order_by_layer) { if (!placement.matches(cfg)) { if (err) *err = "placement does not match config"; return false; @@ -203,6 +296,11 @@ bool build_moe_hybrid_storage(const MoeHybridConfig & cfg, if (err) *err = "layer_descs size does not match n_layer"; return false; } + if (cold_expert_order_by_layer && + (int)cold_expert_order_by_layer->size() != cfg.n_layer) { + if (err) *err = "cold expert override does not match n_layer"; + return false; + } out.placement = placement; out.layers.resize((size_t)cfg.n_layer); @@ -223,31 +321,15 @@ bool build_moe_hybrid_storage(const MoeHybridConfig & cfg, } dst.hot_expert_ids = placement.hot_expert_ids[(size_t)il]; - dst.hot_local_by_global.assign((size_t)cfg.n_expert, -1); - dst.cold_local_by_global.assign((size_t)cfg.n_expert, -1); - - std::vector is_hot((size_t)cfg.n_expert, 0); - for (size_t i = 0; i < dst.hot_expert_ids.size(); ++i) { - const int32_t expert = dst.hot_expert_ids[i]; - if (expert < 0 || expert >= cfg.n_expert) { - if (err) *err = "hot expert id out of range"; - return false; - } - dst.hot_local_by_global[(size_t)expert] = (int32_t)i; - is_hot[(size_t)expert] = 1; - } - for (int expert = 0; expert < cfg.n_expert; ++expert) { - if (!is_hot[(size_t)expert]) { - dst.cold_local_by_global[(size_t)expert] = (int32_t)dst.cold_expert_ids.size(); - dst.cold_expert_ids.push_back((int32_t)expert); - } - } - - // Populate VRAM bitmask from hot expert IDs - std::memset(dst.expert_vram_mask, 0, sizeof(dst.expert_vram_mask)); - for (int32_t eid : dst.hot_expert_ids) { - if (eid >= 0 && eid < 256) - dst.expert_vram_mask[eid >> 6] |= (1ULL << (eid & 63)); + const std::vector * cold_override = + cold_expert_order_by_layer + ? &(*cold_expert_order_by_layer)[(size_t)il] + : nullptr; + if (!build_layer_expert_residency( + cfg, dst.hot_expert_ids, cold_override, + dst.hot_local_by_global, dst.cold_local_by_global, + dst.cold_expert_ids, dst.expert_vram_mask, err)) { + return false; } dst.fused_gate_up = desc.has_fused_gate_up(); @@ -375,7 +457,9 @@ bool build_moe_hybrid_storage_from_file( const std::vector & file_data, MoeHybridStorage & out, std::string * err, - int cache_slots) { + int cache_slots, + bool load_cold_tensors, + const std::vector> * cold_expert_order_by_layer) { if (!placement.matches(cfg)) { if (err) *err = "placement does not match config"; @@ -385,6 +469,11 @@ bool build_moe_hybrid_storage_from_file( if (err) *err = "layer_descs/file_data size does not match n_layer"; return false; } + if (cold_expert_order_by_layer && + (int)cold_expert_order_by_layer->size() != cfg.n_layer) { + if (err) *err = "cold expert override does not match n_layer"; + return false; + } out.placement = placement; out.layers.resize((size_t)cfg.n_layer); @@ -406,31 +495,15 @@ bool build_moe_hybrid_storage_from_file( } dst.hot_expert_ids = placement.hot_expert_ids[(size_t)il]; - dst.hot_local_by_global.assign((size_t)cfg.n_expert, -1); - dst.cold_local_by_global.assign((size_t)cfg.n_expert, -1); - - std::vector is_hot((size_t)cfg.n_expert, 0); - for (size_t i = 0; i < dst.hot_expert_ids.size(); ++i) { - const int32_t expert = dst.hot_expert_ids[i]; - if (expert < 0 || expert >= cfg.n_expert) { - if (err) *err = "hot expert id out of range"; - return false; - } - dst.hot_local_by_global[(size_t)expert] = (int32_t)i; - is_hot[(size_t)expert] = 1; - } - for (int expert = 0; expert < cfg.n_expert; ++expert) { - if (!is_hot[(size_t)expert]) { - dst.cold_local_by_global[(size_t)expert] = (int32_t)dst.cold_expert_ids.size(); - dst.cold_expert_ids.push_back((int32_t)expert); - } - } - - // Populate VRAM bitmask from hot expert IDs - std::memset(dst.expert_vram_mask, 0, sizeof(dst.expert_vram_mask)); - for (int32_t eid : dst.hot_expert_ids) { - if (eid >= 0 && eid < 256) - dst.expert_vram_mask[eid >> 6] |= (1ULL << (eid & 63)); + const std::vector * cold_override = + cold_expert_order_by_layer + ? &(*cold_expert_order_by_layer)[(size_t)il] + : nullptr; + if (!build_layer_expert_residency( + cfg, dst.hot_expert_ids, cold_override, + dst.hot_local_by_global, dst.cold_local_by_global, + dst.cold_expert_ids, dst.expert_vram_mask, err)) { + return false; } dst.fused_gate_up = desc.has_fused_gate_up(); @@ -506,7 +579,7 @@ bool build_moe_hybrid_storage_from_file( } // Allocate cold expert tensors on CPU - if (cold_count > 0) { + if (cold_count > 0 && load_cold_tensors) { ggml_init_params ip{}; ip.mem_size = 16 * ggml_tensor_overhead(); ip.mem_buffer = nullptr; @@ -651,10 +724,14 @@ bool build_moe_hybrid_storage_from_file_with_mmap( size_t mmap_total_size, MoeHybridStorage & out, std::string * err, - int cache_slots) { + int cache_slots, + const std::vector> * cold_expert_order_by_layer) { // First build storage normally (hot GPU + cold CPU buffers). - if (!build_moe_hybrid_storage_from_file(cfg, gpu_backend, placement, layer_descs, file_data, out, err, cache_slots)) { + if (!build_moe_hybrid_storage_from_file( + cfg, gpu_backend, placement, layer_descs, file_data, out, + err, cache_slots, /*load_cold_tensors=*/true, + cold_expert_order_by_layer)) { return false; } diff --git a/server/src/common/moe_hybrid_storage.h b/server/src/common/moe_hybrid_storage.h index d4a1d47d4..7b9eec5ef 100644 --- a/server/src/common/moe_hybrid_storage.h +++ b/server/src/common/moe_hybrid_storage.h @@ -64,7 +64,13 @@ struct CachedHotBatchedGraph { ggml_tensor * sel = nullptr; // [n_used, n_tokens] I32 hot-local IDs ggml_tensor * wts = nullptr; // [n_used, n_tokens] F32 expert weights ggml_tensor * output = nullptr; // [n_embd, n_tokens] F32 output + ggml_tensor * global_ids = nullptr; // [n_used, n_tokens] I32 global expert ids + ggml_tensor * raw_weights = nullptr; // [n_used, n_tokens] F32 router weights + ggml_tensor * hot_local_lut = nullptr; // [1,n_expert] I32 global->local hot id + ggml_tensor * valid_lut = nullptr; // [1,n_expert] F32 1=hot 0=cold + ggml_tensor * residual_in = nullptr; // [n_embd, n_tokens] F32 residual (gpu-remap) int n_tokens = 0; + bool gpu_remap = false; bool valid() const { return ctx && gf && alloc && output; } void free(); @@ -100,7 +106,8 @@ struct MoeHybridLayerStorage { std::vector spare_lru; // [cache_slots] last-use tick uint64_t lru_clock = 0; - // Bitmask: bit set = expert is in VRAM (hot). Supports up to 256 experts. + // Bitmask: bit set = expert is in VRAM (hot). Fast-path cache for the first + // 256 expert ids; correctness must still come from hot_local_by_global. uint64_t expert_vram_mask[4] = {}; // Fast check: are ALL routed experts in VRAM for this batch? @@ -108,9 +115,12 @@ struct MoeHybridLayerStorage { bool all_routed_are_hot(const int32_t * selected_ids, int n_slots) const { for (int i = 0; i < n_slots; ++i) { const int g = selected_ids[i]; - if (g < 0 || g >= 256) continue; - if (!((expert_vram_mask[g >> 6] >> (g & 63)) & 1ULL)) + if (g < 0 || g >= (int)hot_local_by_global.size()) { return false; + } + if (hot_local_by_global[(size_t)g] < 0) { + return false; + } } return true; } @@ -192,7 +202,8 @@ bool build_moe_hybrid_storage(const MoeHybridConfig & cfg, const MoeHybridPlacement & placement, const std::vector & layer_descs, MoeHybridStorage & out, - std::string * err = nullptr); + std::string * err = nullptr, + const std::vector> * cold_expert_order_by_layer = nullptr); // Swap a cold expert into a spare GPU cache slot (LRU evict). Returns the new // hot-local index, or -1 on failure. No-op (returns existing) if already hot. @@ -208,7 +219,9 @@ bool build_moe_hybrid_storage_from_file( const std::vector & file_data, MoeHybridStorage & out, std::string * err = nullptr, - int cache_slots = 0); + int cache_slots = 0, + bool load_cold_tensors = true, + const std::vector> * cold_expert_order_by_layer = nullptr); // Spark: split a VRAM budget into a pinned-hot tier + an auto-sized expert // cache ring. target_bytes==0 keeps the current budget (use the card); @@ -233,6 +246,7 @@ bool build_moe_hybrid_storage_from_file_with_mmap( size_t mmap_total_size, MoeHybridStorage & out, std::string * err = nullptr, - int cache_slots = 0); + int cache_slots = 0, + const std::vector> * cold_expert_order_by_layer = nullptr); } // namespace dflash::common diff --git a/server/src/internal.h b/server/src/internal.h index 89b0f162e..290cfc66d 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -160,6 +160,10 @@ struct TargetWeights { // Metadata from GGUF (validated at load time) int full_attention_interval = 4; + // Per-layer structural type derived from tensor presence. This preserves + // the historical modulo pattern for older qwen35 models while supporting + // layouts such as 49-layer qwen35moe with consecutive tail attention. + std::vector full_attention_layers; int rope_sections[4] = {11, 11, 10, 0}; int n_embd_head_k = 256; // key_length int n_embd_head_v = 256; // value_length @@ -203,6 +207,46 @@ struct TargetWeights { int capture_layer_ids[DFLASH27B_DRAFT_N_TARGET_LAYERS] = {1, 16, 31, 46, 61}; }; +inline bool qwen35_layer_is_full_attention(const TargetWeights & w, int layer_idx) { + if (layer_idx < 0 || layer_idx >= w.n_layer) return false; + if ((int)w.full_attention_layers.size() == w.n_layer) { + return w.full_attention_layers[(size_t)layer_idx] != 0; + } + return w.full_attention_interval > 0 && + (((layer_idx + 1) % w.full_attention_interval) == 0); +} + +inline int qwen35_count_full_attention_layers(const TargetWeights & w) { + if ((int)w.full_attention_layers.size() == w.n_layer) { + int count = 0; + for (uint8_t v : w.full_attention_layers) { + if (v) ++count; + } + return count; + } + return w.full_attention_interval > 0 ? w.n_layer / w.full_attention_interval : 0; +} + +inline int qwen35_count_deltanet_layers(const TargetWeights & w) { + return w.n_layer - qwen35_count_full_attention_layers(w); +} + +inline int qwen35_full_attention_index(const TargetWeights & w, int layer_idx) { + int index = 0; + for (int il = 0; il < layer_idx; ++il) { + if (qwen35_layer_is_full_attention(w, il)) ++index; + } + return index; +} + +inline int qwen35_deltanet_index(const TargetWeights & w, int layer_idx) { + int index = 0; + for (int il = 0; il < layer_idx; ++il) { + if (!qwen35_layer_is_full_attention(w, il)) ++index; + } + return index; +} + // Check if a token is an end-of-sequence marker for the given target weights. inline bool is_eos_tok(int tok, const TargetWeights & w) { return (w.eos_chat_id >= 0 && tok == w.eos_chat_id) diff --git a/server/src/ipc/backend_ipc_main.cpp b/server/src/ipc/backend_ipc_main.cpp index 8f4caf749..854d55f2a 100644 --- a/server/src/ipc/backend_ipc_main.cpp +++ b/server/src/ipc/backend_ipc_main.cpp @@ -1,6 +1,7 @@ // Standalone backend IPC daemon entry point. #include "backend_ipc.h" +#include "moe_expert_compute.h" #include "dflash_draft_ipc.h" #include "gemma4/gemma4_layer_split_adapter.h" #include "laguna/laguna_layer_split_adapter.h" @@ -123,7 +124,10 @@ int main(int argc, char ** argv) { " or: %s --backend-ipc-mode=laguna-target-shard " "--stream-fd=FD --target-gpus=N[,N...] --layer-begins=N[,N...] " "--layer-ends=N[,N...] --max-ctx=N " - "[--hidden=N --vocab=N --max-tokens=N]\n", + "[--hidden=N --vocab=N --max-tokens=N]\n" + " or: %s --backend-ipc-mode=moe-expert-compute " + "--stream-fd=FD --target-gpu=N --placement=PATH\n", + argv[0], argv[0], argv[0], argv[0], @@ -153,6 +157,7 @@ int main(int argc, char ** argv) { size_t shared_payload_bytes = 0; bool enable_dflash = false; int kvflash_pool_tokens = 0; + const char * placement_path = nullptr; for (int i = arg_begin; i < argc; i++) { if (std::strncmp(argv[i], "--ring-cap=", 11) == 0) { if (!parse_nonnegative_int(argv[i] + 11, ring_cap)) return 2; @@ -274,6 +279,10 @@ int main(int argc, char ** argv) { const char * value = nullptr; if (!require_value(i, argc, argv, "--kvflash-pool", value)) return 2; if (!parse_nonnegative_int(value, kvflash_pool_tokens)) return 2; + } else if (std::strncmp(argv[i], "--placement=", 12) == 0) { + placement_path = argv[i] + 12; + } else if (std::strcmp(argv[i], "--placement") == 0) { + if (!require_value(i, argc, argv, "--placement", placement_path)) return 2; } else { std::fprintf(stderr, "[backend-ipc-daemon] unknown option: %s\n", argv[i]); return 2; @@ -318,6 +327,11 @@ int main(int argc, char ** argv) { payload_path, target_gpus, layer_begins, layer_ends, max_ctx, stream_fd, payload_fd, shared_payload_fd, shared_payload_bytes, kvflash_pool_tokens); + case BackendIpcMode::MoeExpertCompute: + return run_moe_expert_compute_ipc_daemon(payload_path, placement_path, + target_gpu, stream_fd, + payload_fd, shared_payload_fd, + shared_payload_bytes); } std::fprintf(stderr, "[backend-ipc-daemon] unsupported mode\n"); return 2; diff --git a/server/src/qwen35/gguf_target_loader.cpp b/server/src/qwen35/gguf_target_loader.cpp index 21788d8f7..8a59f85ce 100644 --- a/server/src/qwen35/gguf_target_loader.cpp +++ b/server/src/qwen35/gguf_target_loader.cpp @@ -45,9 +45,8 @@ #include "internal.h" #include "common/derived_scalars.h" +#include "common/gguf_tensor_data.h" #include "common/layer_split_utils.h" -#include "common/gguf_mmap.h" -#include "common/gguf_bounds.h" #include #include @@ -58,10 +57,7 @@ #include #if !defined(_WIN32) -#include -#include #include -#include #include #endif @@ -177,8 +173,6 @@ static bool should_load_target_tensor(const char * name, struct TargetTensorAlloc { ggml_tensor * tensor = nullptr; - size_t file_offset = 0; - size_t file_size = 0; size_t buffer_offset = 0; }; @@ -202,13 +196,16 @@ bool load_target_gguf_partial(const std::string & path, TargetWeights & out) { // ── 1. Parse metadata + create a ggml_context holding tensor descriptors ─ - ggml_context * meta_ctx = nullptr; - gguf_init_params gip{}; - gip.no_alloc = true; - gip.ctx = &meta_ctx; - gguf_context * gctx = gguf_init_from_file(path.c_str(), gip); - if (!gctx) { - set_last_error("gguf_init_from_file failed: " + path); + std::string err; + GgufTensorDataReader gguf_reader; + if (!gguf_reader.open(path, /*build_merged_tensor_context=*/true, err)) { + set_last_error(err); + return false; + } + ggml_context * meta_ctx = gguf_reader.merged_context(); + const gguf_context * gctx = gguf_reader.primary_context(); + if (!meta_ctx || !gctx) { + set_last_error("GGUF tensor reader did not initialize metadata"); return false; } @@ -219,7 +216,6 @@ bool load_target_gguf_partial(const std::string & path, int64_t arch_id = gguf_find_key(gctx, "general.architecture"); if (arch_id < 0) { set_last_error("missing general.architecture"); - gguf_free(gctx); return false; } const char * arch = gguf_get_val_str(gctx, arch_id); @@ -228,12 +224,10 @@ bool load_target_gguf_partial(const std::string & path, if (arch_str != "qwen35" && arch_str != "qwen35moe") { set_last_error(std::string("unexpected arch: ") + arch_str + " (expected qwen35 or qwen35moe)"); - gguf_free(gctx); return false; } } - std::string err; auto key = [&](const char * suffix) { return arch_str + "." + suffix; }; @@ -279,20 +273,21 @@ bool load_target_gguf_partial(const std::string & path, n_ff_exp, n_ff_shexp, n_expert, n_expert_used, fai, ssm_conv, ssm_inner, ssm_state, ssm_dt, ssm_grp); set_last_error(buf); - gguf_free(gctx); return false; } // Structural invariants required by the graph builder. if (kl != vl) { set_last_error("key_length != value_length not supported"); - gguf_free(gctx); return false; + return false; } - if (n_layer % fai != 0) { + if (fai > n_layer) { char buf[128]; - std::snprintf(buf, sizeof(buf), "block_count=%u not divisible by full_attention_interval=%u", n_layer, fai); + std::snprintf(buf, sizeof(buf), + "full_attention_interval=%u exceeds block_count=%u", + fai, n_layer); set_last_error(buf); - gguf_free(gctx); return false; + return false; } // rope dimension_sections (array of 4 uint32) @@ -302,12 +297,12 @@ bool load_target_gguf_partial(const std::string & path, int64_t rid = gguf_find_key(gctx, rope_sections_key.c_str()); if (rid < 0) { set_last_error("missing rope.dimension_sections"); - gguf_free(gctx); return false; + return false; } size_t n = gguf_get_arr_n(gctx, rid); if (n < 4) { set_last_error("qwen35.rope.dimension_sections has < 4 entries"); - gguf_free(gctx); return false; + return false; } const int32_t * arr = (const int32_t *)gguf_get_arr_data(gctx, rid); for (int k = 0; k < 4; k++) rope_sections[k] = arr[k]; @@ -324,7 +319,7 @@ bool load_target_gguf_partial(const std::string & path, std::snprintf(buf, sizeof(buf), "rope_sections[%d]=%d is negative", k, rope_sections[k]); set_last_error(buf); - gguf_free(gctx); return false; + return false; } sum += rope_sections[k]; } @@ -336,7 +331,7 @@ bool load_target_gguf_partial(const std::string & path, rope_sections[0], rope_sections[1], rope_sections[2], rope_sections[3], n_rot, kl); set_last_error(buf); - gguf_free(gctx); return false; + return false; } } @@ -350,11 +345,10 @@ bool load_target_gguf_partial(const std::string & path, "invalid target load layer range [%d,%d) for n_layer=%u", plan.layer_begin, plan.layer_end, n_layer); set_last_error(buf); - gguf_free(gctx); return false; } - out.ctx = meta_ctx; + out.ctx = nullptr; out.backend = backend; out.n_layer = (int)n_layer; out.n_embd = (int)n_embd; @@ -403,6 +397,7 @@ bool load_target_gguf_partial(const std::string & path, } out.layers.assign((size_t)n_layer, TargetLayer{}); + out.full_attention_layers.assign((size_t)n_layer, 0); // ── 2. Wire our layer pointers to tensors inside meta_ctx ───────── auto g = [&](const char * name) -> ggml_tensor * { @@ -413,7 +408,6 @@ bool load_target_gguf_partial(const std::string & path, out.output = g("output.weight"); if (!out.tok_embd || !out.out_norm || !out.output) { set_last_error("missing top-level tensors (token_embd/output_norm/output)"); - gguf_free(gctx); return false; } out.n_vocab = (int)out.tok_embd->ne[1]; @@ -433,7 +427,6 @@ bool load_target_gguf_partial(const std::string & path, char b[128]; std::snprintf(b, sizeof(b), "layer %d: missing shared norm tensor", il); set_last_error(b); - gguf_free(gctx); return false; } if (is_moe) { @@ -454,7 +447,6 @@ bool load_target_gguf_partial(const std::string & path, char b[128]; std::snprintf(b, sizeof(b), "layer %d: missing dense FFN tensor", il); set_last_error(b); - gguf_free(gctx); return false; } } @@ -482,23 +474,18 @@ bool load_target_gguf_partial(const std::string & path, // NVFP4 per-tensor weight scales are read after the mmap is loaded (below). // Sanity: each layer must be EITHER full-attn OR deltanet, not both, not neither. + // The common qwen35 pattern is modulo-based, but larger qwen35moe + // GGUFs can have special tail layers. Use actual tensor presence as + // the source of truth for downstream cache and graph layout. const bool has_attn = L.wq && L.wk && L.wv && L.wo && L.q_norm && L.k_norm; const bool has_ssm = L.wqkv && L.wqkv_gate && L.ssm_conv1d && L.ssm_out; - const bool is_full_attn_layer = (((il + 1) % out.full_attention_interval) == 0); - if (is_full_attn_layer && !has_attn) { - char b[128]; - std::snprintf(b, sizeof(b), "layer %d expected full-attn, missing tensors", il); - set_last_error(b); - gguf_free(gctx); - return false; - } - if (!is_full_attn_layer && !has_ssm) { + if (has_attn == has_ssm) { char b[128]; - std::snprintf(b, sizeof(b), "layer %d expected deltanet, missing tensors", il); + std::snprintf(b, sizeof(b), "layer %d has ambiguous attention/ssm tensors", il); set_last_error(b); - gguf_free(gctx); return false; } + out.full_attention_layers[(size_t)il] = has_attn ? 1 : 0; if (is_moe) { const bool has_routed = L.ffn_gate_inp && L.ffn_down_exps && @@ -511,7 +498,6 @@ bool load_target_gguf_partial(const std::string & path, char b[160]; std::snprintf(b, sizeof(b), "layer %d expected moe FFN tensors", il); set_last_error(b); - gguf_free(gctx); return false; } } @@ -522,32 +508,33 @@ bool load_target_gguf_partial(const std::string & path, const size_t alignment = ggml_backend_buft_get_alignment(buft); std::vector allocs; size_t alloc_total = 0; - const int64_t n_tensors = gguf_get_n_tensors(gctx); - for (int64_t tid = 0; tid < n_tensors; tid++) { - const char * tname = gguf_get_tensor_name(gctx, tid); - ggml_tensor * t = ggml_get_tensor(meta_ctx, tname); - if (!t || !should_load_target_tensor(tname, plan.layer_begin, plan.layer_end, plan.load_output, plan.skip_expert_tensors)) { - continue; + for (int si = 0; si < gguf_reader.shard_count(); ++si) { + const gguf_context * shard_gctx = gguf_reader.shard_context(si); + const int64_t n_tensors = gguf_get_n_tensors(shard_gctx); + for (int64_t tid = 0; tid < n_tensors; tid++) { + const char * tname = gguf_get_tensor_name(shard_gctx, tid); + ggml_tensor * t = ggml_get_tensor(meta_ctx, tname); + if (!t || !should_load_target_tensor(tname, plan.layer_begin, + plan.layer_end, plan.load_output, + plan.skip_expert_tensors)) { + continue; + } + alloc_total = align_up_size(alloc_total, alignment); + TargetTensorAlloc a; + a.tensor = t; + a.buffer_offset = alloc_total; + alloc_total += ggml_backend_buft_get_alloc_size(buft, t); + allocs.push_back(a); } - alloc_total = align_up_size(alloc_total, alignment); - TargetTensorAlloc a; - a.tensor = t; - a.file_offset = gguf_get_data_offset(gctx) + gguf_get_tensor_offset(gctx, tid); - a.file_size = gguf_get_tensor_size(gctx, tid); - a.buffer_offset = alloc_total; - alloc_total += ggml_backend_buft_get_alloc_size(buft, t); - allocs.push_back(a); } if (allocs.empty()) { set_last_error("target load plan selected no GPU tensors"); - gguf_free(gctx); return false; } out.buf = ggml_backend_alloc_buffer(backend, alloc_total); if (!out.buf) { set_last_error("ggml_backend_alloc_ctx_tensors failed (target)"); - gguf_free(gctx); return false; } ggml_backend_buffer_set_usage(out.buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); @@ -556,52 +543,49 @@ bool load_target_gguf_partial(const std::string & path, for (const TargetTensorAlloc & a : allocs) { if (ggml_backend_tensor_alloc(out.buf, a.tensor, base + a.buffer_offset) != GGML_STATUS_SUCCESS) { set_last_error("ggml_backend_tensor_alloc failed (target)"); - gguf_free(gctx); return false; } } - // ── 4. mmap the file and copy tensor bytes to CUDA ──────────────── + // ── 4. mmap the file(s) and copy tensor bytes to CUDA ───────────── // // SKIP uploading token_embd.weight — it stays on CPU for embedding - // lookup (CUDA get_rows doesn't support k-quants). Its bytes are copied - // into owned host memory below (step 5), so the mmap is released when this - // local goes out of scope. - GgufMmap mm; - if (!mm.open(path, err)) { set_last_error(err); gguf_free(gctx); return false; } - const uint8_t * mm_addr = (const uint8_t *)mm.data(); - const size_t mm_len = mm.size(); - const size_t data_start = gguf_get_data_offset(gctx); + // lookup (CUDA get_rows doesn't support k-quants). Split GGUF shards are + // supported by resolving each tensor to its owning shard mmap. + if (!gguf_reader.open_mmaps(err)) { + set_last_error(err); + return false; + } size_t total = 0; - size_t tok_embd_off = 0, tok_embd_sz = 0; + const uint8_t * tok_embd_data = nullptr; + size_t tok_embd_sz = 0; ggml_type tok_embd_type = GGML_TYPE_COUNT; - for (int64_t tid = 0; tid < n_tensors; tid++) { - const char * tname = gguf_get_tensor_name(gctx, tid); - ggml_tensor * t = ggml_get_tensor(meta_ctx, tname); - if (!t) continue; - const size_t rel_off = gguf_get_tensor_offset(gctx, tid); - const size_t off = data_start + rel_off; - const size_t sz = gguf_get_tensor_size(gctx, tid); - if (!gguf_tensor_in_file(data_start, rel_off, sz, mm_len)) { - set_last_error(gguf_bounds_error("target GGUF", tname, - ggml_type_name(gguf_get_tensor_type(gctx, tid)), - data_start, rel_off, sz, mm_len)); - gguf_free(gctx); - return false; - } - if (std::string(tname) == "token_embd.weight") { - // Remember offset + size for the CPU embedder; don't upload to GPU. - tok_embd_off = off; - tok_embd_sz = sz; - tok_embd_type = gguf_get_tensor_type(gctx, tid); - continue; - } - if (!should_load_target_tensor(tname, plan.layer_begin, plan.layer_end, plan.load_output, plan.skip_expert_tensors)) { - continue; + for (int si = 0; si < gguf_reader.shard_count(); ++si) { + const gguf_context * shard_gctx = gguf_reader.shard_context(si); + const int64_t n_tensors = gguf_get_n_tensors(shard_gctx); + for (int64_t tid = 0; tid < n_tensors; tid++) { + const char * tname = gguf_get_tensor_name(shard_gctx, tid); + GgufTensorRef ref; + if (!gguf_reader.find_tensor(tname, ref) || !ref.data) { + set_last_error(std::string("tensor '") + tname + "' has no mapped data"); + return false; + } + ggml_tensor * t = ggml_get_tensor(meta_ctx, tname); + if (!t) continue; + if (std::string(tname) == "token_embd.weight") { + tok_embd_data = ref.data; + tok_embd_sz = ref.size; + tok_embd_type = ref.type; + continue; + } + if (!should_load_target_tensor(tname, plan.layer_begin, plan.layer_end, + plan.load_output, plan.skip_expert_tensors)) { + continue; + } + ggml_backend_tensor_set(t, ref.data, 0, ref.size); + total += ref.size; } - ggml_backend_tensor_set(t, mm_addr + off, 0, sz); - total += sz; } // ── 4b. Read NVFP4 per-tensor weight scales (optional; 1.0 for non-NVFP4). @@ -618,18 +602,17 @@ bool load_target_gguf_partial(const std::string & path, { auto read_scale = [&](int il, const char * base) -> float { char sname[128]; + GgufTensorRef ref; // Try "base.scale" first (LibertAI), then "base.weight.scale" (heretic) std::snprintf(sname, sizeof(sname), "blk.%d.%s.scale", il, base); - int64_t stid = gguf_find_tensor(gctx, sname); - if (stid < 0) { + bool found = gguf_reader.find_tensor(sname, ref); + if (!found) { std::snprintf(sname, sizeof(sname), "blk.%d.%s.weight.scale", il, base); - stid = gguf_find_tensor(gctx, sname); + found = gguf_reader.find_tensor(sname, ref); } - if (stid < 0) return 1.0f; - const size_t srel = gguf_get_tensor_offset(gctx, stid); - if (!gguf_tensor_in_file(data_start, srel, sizeof(float), mm_len)) return 1.0f; + if (!found || !ref.data || ref.size < sizeof(float)) return 1.0f; float val; - std::memcpy(&val, mm_addr + data_start + srel, sizeof(float)); + std::memcpy(&val, ref.data, sizeof(float)); return val; }; @@ -674,15 +657,23 @@ bool load_target_gguf_partial(const std::string & path, } } - gguf_free(gctx); - // Structural defense: derive head_dim / n_head / n_head_kv from weight // tensor shapes and assert against GGUF-declared metadata. // Uses the first full-attention layer; deltanet layers don't carry wq/wk. // wq packs Q+gate: ne[1] = n_head * n_embd_head_k * 2. // wk: ne[1] = n_head_kv * n_embd_head_k. wq: ne[0] = n_embd. { - const int fa_il = out.full_attention_interval - 1; + int fa_il = -1; + for (int il = 0; il < out.n_layer; ++il) { + if (qwen35_layer_is_full_attention(out, il)) { + fa_il = il; + break; + } + } + if (fa_il < 0) { + set_last_error("no full-attention layer found"); + return false; + } const TargetLayer & fa = out.layers[(size_t)fa_il]; if (fa.wq && fa.wk) { const int64_t exp_q_dim = (int64_t)out.n_head * out.n_embd_head_k * 2; @@ -701,7 +692,7 @@ bool load_target_gguf_partial(const std::string & path, } } - if (tok_embd_off == 0 || tok_embd_type == GGML_TYPE_COUNT) { + if (!tok_embd_data || tok_embd_sz == 0 || tok_embd_type == GGML_TYPE_COUNT) { set_last_error("token_embd.weight not found or invalid type"); return false; } @@ -713,9 +704,7 @@ bool load_target_gguf_partial(const std::string & path, return false; } out.embedder.tok_embd_owned.resize(tok_embd_sz); - std::memcpy(out.embedder.tok_embd_owned.data(), - mm_addr + tok_embd_off, - tok_embd_sz); + std::memcpy(out.embedder.tok_embd_owned.data(), tok_embd_data, tok_embd_sz); out.embedder.tok_embd_bytes = out.embedder.tok_embd_owned.data(); out.embedder.tok_embd_type = tok_embd_type; out.embedder.n_embd = out.n_embd; @@ -731,6 +720,7 @@ bool load_target_gguf_partial(const std::string & path, tok_embd_sz / (1024.0 * 1024.0), ggml_type_name(tok_embd_type)); set_last_error(summary); + out.ctx = gguf_reader.release_merged_context(); return true; } diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 842314b0d..489aa000f 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -28,7 +28,7 @@ bool build_layer_step( if (kvflash) with_mask = true; step_graph_free(sg); - const bool is_attn = (((layer_idx + 1) % w.full_attention_interval) == 0); + const bool is_attn = qwen35_layer_is_full_attention(w, layer_idx); ggml_init_params ip{}; ip.mem_size = 512 * 1024 * 1024; @@ -122,7 +122,7 @@ bool build_layer_prefn_step( ggml_set_name(sg.inp_embed, "inp_embed"); ggml_set_input(sg.inp_embed); - const bool is_attn = (((layer_idx + 1) % w.full_attention_interval) == 0); + const bool is_attn = qwen35_layer_is_full_attention(w, layer_idx); if (is_attn) { sg.positions = ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 4 * n_tokens); ggml_set_name(sg.positions, "positions"); @@ -211,7 +211,7 @@ bool build_hybrid_full_layer_step( ggml_set_name(sg.inp_embed, "inp_embed"); ggml_set_input(sg.inp_embed); - const bool is_attn = (((layer_idx + 1) % w.full_attention_interval) == 0); + const bool is_attn = qwen35_layer_is_full_attention(w, layer_idx); if (is_attn) { sg.positions = ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 4 * n_tokens); ggml_set_name(sg.positions, "positions"); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 4c9a727da..fbdb0a2b0 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -489,8 +489,8 @@ bool Qwen35Backend::snapshot_adopt(int slot, ggml_context * ctx, auto & snap = prefix_snapshots_[slot]; // Count expected tensor layout from weights. - const int n_full_attn = w_.n_layer / w_.full_attention_interval; - const int n_delta = w_.n_layer - n_full_attn; + const int n_full_attn = qwen35_count_full_attention_layers(w_); + const int n_delta = qwen35_count_deltanet_layers(w_); snap.attn_k_snap.assign(n_full_attn, nullptr); snap.attn_v_snap.assign(n_full_attn, nullptr); diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 22a8e3a61..6c314ab22 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -106,8 +106,8 @@ bool create_target_cache_partial(const TargetWeights & w, max_verify_tokens = DFLASH27B_DRAFT_BLOCK_SIZE; } - const int n_full_attn = w.n_layer / w.full_attention_interval; // 16 - const int n_delta = w.n_layer - n_full_attn; // 48 + const int n_full_attn = qwen35_count_full_attention_layers(w); + const int n_delta = qwen35_count_deltanet_layers(w); const int head_dim = w.n_embd_head_k; const int head_v_dim = w.ssm_d_inner / w.ssm_dt_rank; const int conv_ch = w.ssm_d_inner + 2 * w.ssm_n_group * w.ssm_d_state; @@ -156,7 +156,7 @@ bool create_target_cache_partial(const TargetWeights & w, int fa_idx = 0, dn_idx = 0; for (int il = 0; il < w.n_layer; il++) { - const bool is_attn = (((il + 1) % w.full_attention_interval) == 0); + const bool is_attn = qwen35_layer_is_full_attention(w, il); const bool owns_layer = il >= layer_begin && il < layer_end; if (is_attn) { if (!owns_layer) { fa_idx++; continue; } @@ -227,7 +227,7 @@ bool create_target_cache_partial(const TargetWeights & w, int dn_idx = 0; for (int il = 0; il < w.n_layer; il++) { - if (((il + 1) % w.full_attention_interval) != 0) { + if (!qwen35_layer_is_full_attention(w, il)) { const bool owns_layer = il >= layer_begin && il < layer_end; if (!owns_layer) { dn_idx++; continue; } ggml_tensor * Sn = ggml_new_tensor_3d(out.rollback_ctx, GGML_TYPE_F32, @@ -379,7 +379,7 @@ bool migrate_prefill_cache(const TargetWeights & w, int dn_idx = 0; for (int il = 0; il < w.n_layer; il++) { - if (((il + 1) % w.full_attention_interval) != 0) { + if (!qwen35_layer_is_full_attention(w, il)) { ggml_tensor * Sn = ggml_new_tensor_3d(cache.rollback_ctx, GGML_TYPE_F32, head_v_dim, head_v_dim, w.ssm_dt_rank); ggml_tensor * Cn = ggml_new_tensor_2d(cache.rollback_ctx, GGML_TYPE_F32, @@ -1050,7 +1050,7 @@ static ggml_tensor * build_single_layer( const int hidden = w.n_embd; const float eps = w.rms_eps; const TargetLayer & L = w.layers[layer_idx]; - const bool is_attn = (((layer_idx + 1) % w.full_attention_interval) == 0); + const bool is_attn = qwen35_layer_is_full_attention(w, layer_idx); const int * CAPTURE_LAYERS = w.capture_layer_ids; const int N_CAPTURE = w.n_capture_layers; @@ -1060,10 +1060,7 @@ static ggml_tensor * build_single_layer( ggml_tensor * cur = rms_norm_mul(ctx, inp_f32, L.attn_norm, eps); if (is_attn) { - int fa_idx = 0; - for (int il = 0; il < layer_idx; il++) { - if (((il + 1) % w.full_attention_interval) == 0) fa_idx++; - } + const int fa_idx = qwen35_full_attention_index(w, layer_idx); cur = build_full_attn_block(ctx, gf, w, L, cur, positions, w.rope_sections, cache.attn_k[fa_idx], cache.attn_v[fa_idx], attn_mask, kv_start, n_tokens, @@ -1073,10 +1070,7 @@ static ggml_tensor * build_single_layer( q_tail_capture, q_tail_start, kv_write_rows); } else { - int dn_idx = 0; - for (int il = 0; il < layer_idx; il++) { - if (((il + 1) % w.full_attention_interval) != 0) dn_idx++; - } + const int dn_idx = qwen35_deltanet_index(w, layer_idx); cur = build_delta_net_block(ctx, gf, w, L, cur, cache.conv_state[dn_idx], cache.ssm_state[dn_idx], n_tokens, nullptr, nullptr, @@ -1155,8 +1149,7 @@ QwenGraphOutputs build_qwen35_graph( // net layer count so we can index by dn_idx as we iterate the layers. QwenGraphOutputs og_early{}; if (in.capture_delta_intermediate) { - const int n_full_attn = w.n_layer / w.full_attention_interval; - const int n_delta = w.n_layer - n_full_attn; + const int n_delta = qwen35_count_deltanet_layers(w); og_early.delta_captures.resize(n_delta); } if (in.capture_moe_router && w.is_moe) { @@ -1172,7 +1165,7 @@ QwenGraphOutputs build_qwen35_graph( for (int il = 0; il < w.n_layer; il++) { const TargetLayer & L = w.layers[il]; - const bool is_attn = (((il + 1) % w.full_attention_interval) == 0); + const bool is_attn = qwen35_layer_is_full_attention(w, il); ggml_tensor * inp_f32 = graph_tensor_f32(ctx, inpL); ggml_tensor * inpSA = inp_f32; @@ -1384,17 +1377,14 @@ QwenLayerPrefnOutputs build_qwen35_layer_prefn( QwenLayerPrefnOutputs out{}; const float eps = w.rms_eps; const TargetLayer & L = w.layers[layer_idx]; - const bool is_attn = (((layer_idx + 1) % w.full_attention_interval) == 0); + const bool is_attn = qwen35_layer_is_full_attention(w, layer_idx); ggml_tensor * inp_f32 = graph_tensor_f32(ctx, inp); ggml_tensor * inpSA = inp_f32; ggml_tensor * cur = rms_norm_mul(ctx, inp_f32, L.attn_norm, eps); if (is_attn) { - int fa_idx = 0; - for (int il = 0; il < layer_idx; il++) { - if (((il + 1) % w.full_attention_interval) == 0) fa_idx++; - } + const int fa_idx = qwen35_full_attention_index(w, layer_idx); cur = build_full_attn_block(ctx, gf, w, L, cur, positions, w.rope_sections, cache.attn_k[fa_idx], cache.attn_v[fa_idx], attn_mask, kv_start, n_tokens, @@ -1404,10 +1394,7 @@ QwenLayerPrefnOutputs build_qwen35_layer_prefn( /*q_tail_capture=*/nullptr, /*q_tail_start=*/0, kv_write_rows); } else { - int dn_idx = 0; - for (int il = 0; il < layer_idx; il++) { - if (((il + 1) % w.full_attention_interval) != 0) dn_idx++; - } + const int dn_idx = qwen35_deltanet_index(w, layer_idx); cur = build_delta_net_block(ctx, gf, w, L, cur, cache.conv_state[dn_idx], cache.ssm_state[dn_idx], n_tokens, nullptr, nullptr, @@ -1431,8 +1418,8 @@ bool snapshot_target_cache(const TargetWeights & w, const TargetCache & cache, ggml_backend_t backend, PrefixSnapshot & snap) { - const int n_full_attn = w.n_layer / w.full_attention_interval; // 16 - const int n_delta = w.n_layer - n_full_attn; // 48 + const int n_full_attn = qwen35_count_full_attention_layers(w); + const int n_delta = qwen35_count_deltanet_layers(w); const int snap_pos = cache.cur_pos; if (snap_pos <= 0) { @@ -1676,7 +1663,7 @@ bool snapshot_target_cache_thin(const TargetWeights & w, set_last_error("snapshot_thin: kv_end exceeds cache.cur_pos (would capture uninitialized KV)"); return false; } - const int n_full_attn = w.n_layer / w.full_attention_interval; + const int n_full_attn = qwen35_count_full_attention_layers(w); const int block_size = kv_end - kv_start; // Lazy alloc; if snap was already a THIN with same range, reuse. diff --git a/server/src/qwen35moe/qwen35moe_backend.cpp b/server/src/qwen35moe/qwen35moe_backend.cpp index 9f8559e15..dd25c2fbc 100644 --- a/server/src/qwen35moe/qwen35moe_backend.cpp +++ b/server/src/qwen35moe/qwen35moe_backend.cpp @@ -21,7 +21,9 @@ #include #include #include +#include #include "common/gguf_mmap.h" +#include "common/gguf_tensor_data.h" namespace dflash::common { @@ -33,11 +35,49 @@ static uint64_t elapsed_us(HybridClock::time_point start, HybridClock::time_poin return (uint64_t) std::chrono::duration_cast(end - start).count(); } +static PlacementBackend resolved_device_backend(const DevicePlacement & device) { + return device.backend == PlacementBackend::Auto + ? compiled_placement_backend() + : device.backend; +} + +static bool resolve_expert_split_targets_from_env( + uint64_t primary_capacity_bytes, + std::vector & out, + std::string * err) { + return dflash::common::resolve_expert_split_targets_from_env( + "DFLASH_QWEN35MOE_EXPERT_TARGETS", + "DFLASH_QWEN35MOE_EXPERT_TARGET_CAPS", + primary_capacity_bytes, + out, + err); +} + +static bool same_placement(const MoeHybridPlacement & lhs, + const MoeHybridPlacement & rhs) { + return lhs.n_layer == rhs.n_layer && + lhs.n_expert == rhs.n_expert && + lhs.n_expert_used == rhs.n_expert_used && + lhs.total_hot == rhs.total_hot && + lhs.hot_counts == rhs.hot_counts && + lhs.hot_expert_ids == rhs.hot_expert_ids; +} + } // namespace Qwen35MoeBackend::Qwen35MoeBackend(const Qwen35Config & cfg) : Qwen35Backend(cfg) {} +bool Qwen35MoeBackend::supports_dflash_spec_decode() const { + if (!target_weights().moe_hybrid) { + return true; + } + for (const auto & layer : target_weights().moe_hybrid->layers) { + if (!layer.cold_expert_ids.empty()) return false; + } + return true; +} + bool Qwen35MoeBackend::load_target_model(ggml_backend_t backend, TargetWeights & out) { // Phase 1: Load core model (non-expert tensors) to GPU. // Expert tensors get metadata descriptors but are NOT allocated on GPU. @@ -72,15 +112,11 @@ bool Qwen35MoeBackend::load_target_model(ggml_backend_t backend, TargetWeights & ? std::string("hotness:") + hotness_path : std::string("uniform"); - // If all experts fit on GPU, reload with experts included + // If all experts fit on GPU, keep the hybrid path but note that KVFlash may + // be redundant only when the FULL max_ctx KV reservation still leaves room. if (placement.total_hot >= out.n_layer * out.n_expert) { - std::printf("[qwen35moe] all experts fit in VRAM, loading fully to GPU\n"); + std::printf("[qwen35moe] all experts fit in VRAM, staying on hybrid path (all-hot)\n"); std::fflush(stdout); - // Record the placement result so post_kvflash_init_gate() can disable - // the KVFlash pool (moe_hybrid is null on this all-hot path). - placement_all_hot_ = true; - free_target_weights(out); - return load_target_gguf(cfg_.target_path, backend, out); } if (const char * telemetry = std::getenv("DFLASH_QWEN35MOE_TELEMETRY")) { @@ -88,59 +124,21 @@ bool Qwen35MoeBackend::load_target_model(ggml_backend_t backend, TargetWeights & } // Phase 3: Load expert data from GGUF mmap directly into split hot/cold buffers. - // Open GGUF again to get tensor file offsets and mmap the data. + // Split GGUF shards are supported by resolving each expert tensor to its + // owning shard. Streaming mmap is retained only for single-file GGUFs. { - ggml_context * expert_meta = nullptr; - gguf_init_params gip{}; - gip.no_alloc = true; - gip.ctx = &expert_meta; - gguf_context * gctx = gguf_init_from_file(cfg_.target_path, gip); - if (!gctx) { - set_last_error("failed to re-open GGUF for expert loading"); + GgufTensorDataReader expert_reader; + if (!expert_reader.open(cfg_.target_path, + /*build_merged_tensor_context=*/false, err)) { + set_last_error("failed to open GGUF expert shards: " + err); return false; } - - // Mmap the file - GgufMmap _mf; - std::string _mferr; - if (!_mf.open(cfg_.target_path, _mferr)) { - set_last_error("mmap failed on GGUF: " + _mferr); - gguf_free(gctx); + if (!expert_reader.open_mmaps(err)) { + set_last_error("mmap failed on GGUF expert shards: " + err); return false; } - 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. - 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; - - // Build per-layer expert file data - std::vector layer_file_data((size_t)out.n_layer); - for (int il = 0; il < out.n_layer; ++il) { - char name[128]; - auto find_tensor_data = [&](const char * suffix) -> ExpertTensorFileData { - std::snprintf(name, sizeof(name), "blk.%d.%s.weight", il, suffix); - int64_t tid = gguf_find_tensor(gctx, name); - if (tid < 0) return {}; - size_t off = data_start + gguf_get_tensor_offset(gctx, tid); - size_t sz = gguf_get_tensor_size(gctx, tid); - if (off + sz > file_size) return {}; - return { file_bytes + off, sz }; - }; - - layer_file_data[(size_t)il].gate_exps = find_tensor_data("ffn_gate_exps"); - layer_file_data[(size_t)il].up_exps = find_tensor_data("ffn_up_exps"); - layer_file_data[(size_t)il].down_exps = find_tensor_data("ffn_down_exps"); - layer_file_data[(size_t)il].gate_up_exps = find_tensor_data("ffn_gate_up_exps"); - } + const auto layer_file_data = + make_layer_expert_file_data(expert_reader, out.n_layer); auto hybrid = std::make_shared(); MoeHybridConfig hybrid_cfg = make_moe_hybrid_config(out); @@ -155,20 +153,51 @@ 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 std::vector> * cold_order_by_layer = + last_expert_split_state_.materialization.ordered_cold_union + ? &last_expert_split_state_.materialization.cold_expert_ids_by_layer + : nullptr; + bool storage_ok = false; + if (expert_reader.shard_count() == 1) { + const void * mmap_addr = nullptr; + size_t file_size = 0; + int mmap_fd = -1; + std::string mmap_err; + if (!expert_reader.release_single_mmap(mmap_addr, file_size, + mmap_fd, mmap_err)) { + set_last_error("single GGUF mmap release failed: " + mmap_err); + return false; + } +#if !defined(_WIN32) + if (mmap_fd >= 0) ::close(mmap_fd); +#endif + storage_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, + cold_order_by_layer); + if (!storage_ok) { #if defined(_WIN32) UnmapViewOfFile(const_cast(mmap_addr)); #else ::munmap(const_cast(mmap_addr), file_size); #endif - gguf_free(gctx); - set_last_error(std::string("qwen35moe hybrid storage build failed: ") + err); - return false; } - - // Keep mmap open for streaming prefill — do NOT munmap here. - // The mmap_data/mmap_size are stored in hybrid storage for lifetime management. - gguf_free(gctx); + } else { + storage_ok = build_moe_hybrid_storage_from_file( + hybrid_cfg, backend, placement, layer_descs, layer_file_data, + *hybrid, &err, cache_slots, /*load_cold_tensors=*/true, + cold_order_by_layer); + if (storage_ok) { + std::fprintf(stderr, + "[qwen35moe] split GGUF expert load: %d shards, " + "streaming mmap disabled\n", + expert_reader.shard_count()); + } + } + if (!storage_ok) { + set_last_error(std::string("qwen35moe hybrid storage build failed: ") + err); + return false; + } out.moe_hybrid = std::move(hybrid); } @@ -255,43 +284,25 @@ bool Qwen35MoeBackend::spark_wants_bootstrap() const { // Re-mmap the GGUF and rebuild the hot/cold storage for a new placement. Used by // the Spark bootstrap to apply the calibrated placement in-process. -bool Qwen35MoeBackend::rebuild_hybrid_from_placement(const MoeHybridPlacement & placement, - std::string & err) { +bool Qwen35MoeBackend::rebuild_hybrid_from_placement( + const MoeHybridPlacement & placement, + const ExpertSplitMaterialization * materialization, + std::string & err) { TargetWeights & out = target_weights(); ggml_backend_t backend = target_backend(); - gguf_init_params gip{}; - gguf_context * gctx = gguf_init_from_file(cfg_.target_path, gip); - if (!gctx) { err = "gguf reinit failed"; return false; } - GgufMmap _mf; - std::string _mferr; - if (!_mf.open(cfg_.target_path, _mferr)) { - gguf_free(gctx); - err = "mmap failed: " + _mferr; + GgufTensorDataReader expert_reader; + if (!expert_reader.open(cfg_.target_path, + /*build_merged_tensor_context=*/false, err)) { + err = "gguf expert shard open failed: " + err; return false; } - const size_t file_size = _mf.size(); - const void * mmap_addr = _mf.data(); - - const size_t data_start = gguf_get_data_offset(gctx); - const auto * file_bytes = (const uint8_t *)mmap_addr; - std::vector layer_file_data((size_t)out.n_layer); - for (int il = 0; il < out.n_layer; ++il) { - char name[128]; - auto find_tensor_data = [&](const char * suffix) -> ExpertTensorFileData { - std::snprintf(name, sizeof(name), "blk.%d.%s.weight", il, suffix); - int64_t tid = gguf_find_tensor(gctx, name); - if (tid < 0) return {}; - size_t off = data_start + gguf_get_tensor_offset(gctx, tid); - size_t sz = gguf_get_tensor_size(gctx, tid); - if (off + sz > file_size) return {}; - return { file_bytes + off, sz }; - }; - layer_file_data[(size_t)il].gate_exps = find_tensor_data("ffn_gate_exps"); - layer_file_data[(size_t)il].up_exps = find_tensor_data("ffn_up_exps"); - layer_file_data[(size_t)il].down_exps = find_tensor_data("ffn_down_exps"); - layer_file_data[(size_t)il].gate_up_exps = find_tensor_data("ffn_gate_up_exps"); + if (!expert_reader.open_mmaps(err)) { + err = "mmap failed: " + err; + return false; } + const auto layer_file_data = + make_layer_expert_file_data(expert_reader, out.n_layer); // Free the current hot/cold buffers before allocating the new ones so the // rebuild fits in VRAM (no transient 2x). Safe: bootstrap runs at startup @@ -305,28 +316,155 @@ bool Qwen35MoeBackend::rebuild_hybrid_from_placement(const MoeHybridPlacement & for (int il = 0; il < out.n_layer; ++il) layer_descs[(size_t)il] = make_moe_layer_desc(out.layers[(size_t)il]); const int cache_slots = cache_slots_ >= 0 ? cache_slots_ : 0; + const std::vector> * cold_order_by_layer = nullptr; + if (materialization && materialization->ordered_cold_union) { + if (!materialization->matches(out.n_layer, out.n_expert, out.n_expert_used) || + !same_placement(materialization->primary_placement, placement)) { + err = "expert split materialization does not match rebuild placement"; + return false; + } + cold_order_by_layer = &materialization->cold_expert_ids_by_layer; + } - const bool ok = build_moe_hybrid_storage_from_file(hybrid_cfg, backend, placement, layer_descs, - layer_file_data, *hybrid, &err, cache_slots); - gguf_free(gctx); + const bool ok = build_moe_hybrid_storage_from_file( + hybrid_cfg, backend, placement, layer_descs, layer_file_data, *hybrid, + &err, cache_slots, /*load_cold_tensors=*/true, cold_order_by_layer); if (!ok) return false; out.moe_hybrid = std::move(hybrid); return true; } +bool Qwen35MoeBackend::build_expert_split_plan_from_stats( + const MoeHybridRoutingStats & hotness, + uint64_t expert_budget_bytes, + const TargetWeights & w, + ExpertSplitPlan & out, + std::string * err) const { + if ((int) layer_expert_bytes_.size() != w.n_layer) { + if (err) *err = "layer expert bytes not initialized"; + return false; + } + if (expert_budget_bytes == 0) { + if (err) *err = "expert budget must be > 0"; + return false; + } + + ExpertSplitConfig cfg; + cfg.n_layer = w.n_layer; + cfg.n_expert = w.n_expert; + cfg.allow_implicit_cpu_fallback = true; + cfg.require_full_grid = true; + cfg.min_per_layer_by_target = {std::min(w.n_expert_used, w.n_expert)}; + + std::vector targets; + if (!resolve_expert_split_targets_from_env(expert_budget_bytes, targets, err)) { + return false; + } + if (targets.empty()) { + const PlacementBackend backend = resolved_device_backend(cfg_.device); + targets = { + {std::string(placement_backend_name(backend)) + ":" + std::to_string(cfg_.device.gpu), + placement_backend_name(backend), cfg_.device.gpu, expert_budget_bytes, 0, false}, + }; + } + if (targets.size() > 1) { + cfg.min_per_layer_by_target.assign(targets.size(), 0); + cfg.min_per_layer_by_target[0] = std::min(w.n_expert_used, w.n_expert); + } + + std::vector units; + units.reserve((size_t) w.n_layer * (size_t) w.n_expert); + for (int il = 0; il < w.n_layer; ++il) { + for (int ie = 0; ie < w.n_expert; ++ie) { + units.push_back(ExpertSplitUnit{ + il, + ie, + layer_expert_bytes_[(size_t) il], + (double) hotness.count(il, ie), + }); + } + } + + return build_expert_split_plan(cfg, targets, units, out, err); +} + +bool Qwen35MoeBackend::build_expert_split_state_from_stats( + const MoeHybridRoutingStats & hotness, + uint64_t expert_budget_bytes, + const TargetWeights & w, + std::string * err) { + if ((int) layer_expert_bytes_.size() != w.n_layer) { + if (err) *err = "layer expert bytes not initialized"; + return false; + } + if (expert_budget_bytes == 0) { + if (err) *err = "expert budget must be > 0"; + return false; + } + + ExpertSplitConfig cfg; + cfg.n_layer = w.n_layer; + cfg.n_expert = w.n_expert; + cfg.allow_implicit_cpu_fallback = true; + cfg.require_full_grid = true; + cfg.min_per_layer_by_target = {std::min(w.n_expert_used, w.n_expert)}; + + std::vector targets; + if (!resolve_expert_split_targets_from_env(expert_budget_bytes, targets, err)) { + return false; + } + if (targets.empty()) { + const PlacementBackend backend = resolved_device_backend(cfg_.device); + targets = { + {std::string(placement_backend_name(backend)) + ":" + std::to_string(cfg_.device.gpu), + placement_backend_name(backend), cfg_.device.gpu, expert_budget_bytes, 0, false}, + }; + } + if (targets.size() > 1) { + cfg.min_per_layer_by_target.assign(targets.size(), 0); + cfg.min_per_layer_by_target[0] = std::min(w.n_expert_used, w.n_expert); + } + + std::vector units; + units.reserve((size_t) w.n_layer * (size_t) w.n_expert); + for (int il = 0; il < w.n_layer; ++il) { + for (int ie = 0; ie < w.n_expert; ++ie) { + units.push_back(ExpertSplitUnit{ + il, + ie, + layer_expert_bytes_[(size_t) il], + (double) hotness.count(il, ie), + }); + } + } + + ExpertSplitStateComponents state; + if (!build_expert_split_state(cfg, targets, units, + w.n_expert_used, state, err)) { + return false; + } + if (!validate_primary_expert_split_target(state.plan.targets, + resolved_device_backend(cfg_.device), + cfg_.device.gpu, err)) { + return false; + } + last_expert_split_state_ = std::move(state); + return true; +} + bool Qwen35MoeBackend::spark_bootstrap_finalize(const std::string & profile_path) { if (!spark_wants_bootstrap()) return false; std::string err; routing_stats_->save_csv(profile_path, &err); // persist the observed routing const TargetWeights & w = target_weights(); - MoeHybridPlacement placement; - if (!MoeHybridPlacement::build_from_stats_with_layer_bytes( - *routing_stats_, layer_expert_bytes_, spark_expert_budget_, - std::min(w.n_expert_used, w.n_expert), placement, &err)) { + if (!build_expert_split_state_from_stats(*routing_stats_, spark_expert_budget_, + w, &err)) { std::fprintf(stderr, "[spark] bootstrap placement build failed: %s\n", err.c_str()); return false; } - if (!rebuild_hybrid_from_placement(placement, err)) { + if (!rebuild_hybrid_from_placement( + last_expert_split_state_.materialization.primary_placement, + &last_expert_split_state_.materialization, err)) { std::fprintf(stderr, "[spark] bootstrap storage rebuild failed: %s\n", err.c_str()); return false; } @@ -342,14 +480,14 @@ bool Qwen35MoeBackend::post_kvflash_init_gate() { if (!kvflash_active()) return true; bool should_disable = false; - if (placement_all_hot_full_kv_) { - should_disable = true; - } else if (target_weights().moe_hybrid) { + if (target_weights().moe_hybrid) { int total_cold = 0; for (const auto & ls : target_weights().moe_hybrid->layers) { total_cold += (int)ls.cold_expert_ids.size(); } - if (total_cold == 0) should_disable = true; // hybrid built but 0 cold + if (total_cold == 0 && placement_all_hot_full_kv_) { + should_disable = true; + } } if (should_disable) { @@ -420,8 +558,10 @@ bool Qwen35MoeBackend::run_ar_decode_path(int committed, int n_gen, bool Qwen35MoeBackend::ensure_pipe_state(int kv_start) { if (pipe_state_ && pipe_state_->valid()) return true; pipe_state_ = std::make_unique(); - if (!init_pipelined_decode_state(*pipe_state_, target_backend(), target_weights(), + if (!init_pipelined_decode_state(*pipe_state_, target_backend(), cfg_.target_path, + target_weights(), target_cache(), *target_weights().moe_hybrid, + &last_expert_split_state_.compute_runtime, kv_start, cfg_.kq_stride_pad)) { pipe_state_.reset(); return false; @@ -771,6 +911,12 @@ GenerateResult Qwen35MoeBackend::generate_impl(const GenerateRequest & req, const int prompt_len = (int)req.prompt.size(); const int prefill_chunk = std::min(128, prompt_len); // batch size per GPU compute + if (!ensure_pipe_state(/*kv_start=*/0)) { + result.error = "pipe_state_init"; + cleanup_graphs(); + return result; + } + // kvflash: hybrid prefill writes rows identity-mapped, so the prompt must // fit the pool with one chunk of decode headroom (same contract as the // base do_prefill). @@ -927,7 +1073,9 @@ GenerateResult Qwen35MoeBackend::generate_impl(const GenerateRequest & req, chunk_selected.data(), chunk_weights.data(), chunk_len, ffn_batch_out, &result.error, - &ffn_hot_alloc, &ffn_cold_alloc); + &ffn_hot_alloc, &ffn_cold_alloc, + pipe_state_ ? pipe_state_->moe_expert_compute : nullptr, + pipe_state_ ? pipe_state_->moe_expert_layers.data() + (size_t)il : nullptr); } else if (storage.all_routed_are_hot(chunk_selected.data(), chunk_len * n_expert_used)) { // All selected experts happen to be in VRAM — pure GPU, no CPU @@ -959,7 +1107,9 @@ GenerateResult Qwen35MoeBackend::generate_impl(const GenerateRequest & req, chunk_selected.data(), chunk_weights.data(), chunk_len, ffn_batch_out, &result.error, - &ffn_hot_alloc, &ffn_cold_alloc); + &ffn_hot_alloc, &ffn_cold_alloc, + pipe_state_ ? pipe_state_->moe_expert_compute : nullptr, + pipe_state_ ? pipe_state_->moe_expert_layers.data() + (size_t)il : nullptr); } if (!ffn_ok) { // Per-token fallback (avoids sm_75 mul_mat_id assertion with cold experts) @@ -1760,7 +1910,9 @@ bool Qwen35MoeBackend::hybrid_forward_batch( target_backend(), target_weights().moe_hybrid->cpu_backend, chunk_cfg, chunk_desc, storage, chunk_post.data(), chunk_selected.data(), chunk_weights.data(), - n_tokens, ffn_batch_out, nullptr, &ffn_hot_alloc, nullptr); + n_tokens, ffn_batch_out, nullptr, &ffn_hot_alloc, nullptr, + pipe_state_ ? pipe_state_->moe_expert_compute : nullptr, + pipe_state_ ? pipe_state_->moe_expert_layers.data() + (size_t)il : nullptr); } if (!ffn_ok) { @@ -2268,12 +2420,10 @@ bool Qwen35MoeBackend::load_dynamic_placement(const char * hotness_path, layer_expert_bytes_ = layer_expert_bytes; // Build placement using greedy knapsack with byte budget - if (!MoeHybridPlacement::build_from_stats_with_layer_bytes( - hotness, layer_expert_bytes, expert_budget, - /*min_hot_per_layer=*/std::min(w.n_expert_used, w.n_expert), - out, err)) { + if (!build_expert_split_state_from_stats(hotness, expert_budget, w, err)) { return false; } + out = last_expert_split_state_.materialization.primary_placement; std::printf("[qwen35moe] dynamic placement result: %d hot experts, %d cold experts\n", out.total_hot, w.n_layer * w.n_expert - out.total_hot); diff --git a/server/src/qwen35moe/qwen35moe_backend.h b/server/src/qwen35moe/qwen35moe_backend.h index a731b4f7a..733276cc6 100644 --- a/server/src/qwen35moe/qwen35moe_backend.h +++ b/server/src/qwen35moe/qwen35moe_backend.h @@ -10,6 +10,12 @@ #include "../common/moe_hybrid_stream.h" #include "../common/moe_hybrid_routing_stats.h" #include "../common/moe_hybrid_swap_manager.h" +#include "../common/expert_split_plan.h" +#include "../common/expert_split_runtime.h" +#include "../common/expert_split_compute_runtime.h" +#include "../common/expert_split_materialization.h" +#include "../common/expert_split_state.h" +#include "../common/expert_split_target_config.h" #include "../common/moe_routing_collector.h" #include @@ -27,7 +33,7 @@ class Qwen35MoeBackend : public Qwen35Backend { GenerateResult restore_and_generate_impl(int slot, const GenerateRequest & req, const DaemonIO & io) override; - bool supports_dflash_spec_decode() const override { return !target_weights().moe_hybrid; } + bool supports_dflash_spec_decode() const override; bool set_routing_collector(MoeRoutingCollector * c) override { routing_collector_ = c; return true; } const MoeHybridRoutingStats * get_routing_stats() const override { return routing_stats_.get(); } @@ -44,12 +50,9 @@ class Qwen35MoeBackend : public Qwen35Backend { void after_target_compute(StepGraph & sg, int kv_start, int n_tokens) override; private: - // All-hot placement signal for post_kvflash_init_gate(): set when - // load_target_model takes the all-hot early-return (moe_hybrid null). - bool placement_all_hot_ = false; // True iff all experts fit hot with the FULL max_ctx KV reservation - // (KVFlash redundant). When false but placement_all_hot_ is true, the pool - // is what kept experts hot — the gate must NOT disable KVFlash. + // (KVFlash redundant). When false, the pool is what kept experts hot — + // the gate must NOT disable KVFlash. bool placement_all_hot_full_kv_ = false; std::shared_ptr routing_stats_; std::string routing_stats_out_path_; @@ -57,13 +60,25 @@ class Qwen35MoeBackend : public Qwen35Backend { int cache_slots_ = -1; // Spark auto-sized (-1=unset) uint64_t spark_expert_budget_ = 0; // hot budget, for the bootstrap rebuild std::vector layer_expert_bytes_; - bool rebuild_hybrid_from_placement(const MoeHybridPlacement & placement, std::string & err); + ExpertSplitStateComponents last_expert_split_state_; + bool rebuild_hybrid_from_placement(const MoeHybridPlacement & placement, + const ExpertSplitMaterialization * materialization, + std::string & err); MoeHybridSwapPolicy swap_policy_; bool hybrid_telemetry_ = false; MoeHybridStreamEngine stream_engine_; MoeRoutingCollector * routing_collector_ = nullptr; void maybe_post_request_swap(); + bool build_expert_split_plan_from_stats(const MoeHybridRoutingStats & hotness, + uint64_t expert_budget_bytes, + const TargetWeights & w, + ExpertSplitPlan & out, + std::string * err) const; + bool build_expert_split_state_from_stats(const MoeHybridRoutingStats & hotness, + uint64_t expert_budget_bytes, + const TargetWeights & w, + std::string * err); bool load_dynamic_placement(const char * hotness_path, ggml_backend_t backend, const TargetWeights & w, diff --git a/server/src/qwen35moe/qwen35moe_pipelined_decode.cpp b/server/src/qwen35moe/qwen35moe_pipelined_decode.cpp index 9c1946c0c..6504982fb 100644 --- a/server/src/qwen35moe/qwen35moe_pipelined_decode.cpp +++ b/server/src/qwen35moe/qwen35moe_pipelined_decode.cpp @@ -20,6 +20,66 @@ static uint64_t pipe_elapsed_us(PipelineClock::time_point s, PipelineClock::time return (uint64_t)std::chrono::duration_cast(e - s).count(); } +static bool prepare_unique_decode_hot_slots( + const MoeHybridLayerStorage & storage, + const int32_t * global_ids, + const float * router_weights, + int n_expert_used, + std::vector & hot_ids, + std::vector & hot_weights, + std::vector & cold_ids, + std::vector & cold_weights, + int & n_cold) { + const int n_hot_init = std::max(1, storage.hot_active); + std::vector used((size_t)n_hot_init, 0); + n_cold = 0; + + for (int i = 0; i < n_expert_used; ++i) { + hot_ids[(size_t)i] = -1; + hot_weights[(size_t)i] = 0.0f; + const int32_t gid = global_ids[(size_t)i]; + const int32_t hot_local = + (gid >= 0 && gid < (int)storage.hot_local_by_global.size()) + ? storage.hot_local_by_global[(size_t)gid] + : -1; + if (hot_local >= 0) { + if (hot_local >= n_hot_init || used[(size_t)hot_local]) { + return false; + } + hot_ids[(size_t)i] = hot_local; + hot_weights[(size_t)i] = router_weights[(size_t)i]; + used[(size_t)hot_local] = 1; + continue; + } + if (gid >= 0 && gid < (int)storage.cold_local_by_global.size()) { + const int32_t cold_local = storage.cold_local_by_global[(size_t)gid]; + if (cold_local >= 0 && n_cold < (int)cold_ids.size() && + n_cold < (int)cold_weights.size()) { + cold_ids[(size_t)n_cold] = cold_local; + cold_weights[(size_t)n_cold] = router_weights[(size_t)i]; + ++n_cold; + } + } + } + + int next_dummy = 0; + for (int i = 0; i < n_expert_used; ++i) { + if (hot_ids[(size_t)i] >= 0) continue; + while (next_dummy < n_hot_init && used[(size_t)next_dummy]) { + ++next_dummy; + } + if (next_dummy >= n_hot_init) { + return false; + } + hot_ids[(size_t)i] = next_dummy; + hot_weights[(size_t)i] = 0.0f; + used[(size_t)next_dummy] = 1; + ++next_dummy; + } + + return true; +} + // ─── CachedPrefnGraph ───────────────────────────────────────────────────────── void CachedPrefnGraph::free() { @@ -68,9 +128,7 @@ static bool build_cached_deltanet_prefn( QwenLayerPrefnOutputs go = build_qwen35_layer_prefn( out.ctx, out.gf, w, cache, layer_idx, out.inp_embed, /*positions=*/nullptr, /*attn_mask=*/nullptr, - kv_start, /*n_tokens=*/1, /*fa_window=*/0, - /*kv_write_rows=*/nullptr, - /*skip_gdn_intermediate=*/true); + kv_start, /*n_tokens=*/1, /*fa_window=*/0); if (!go.residual || !go.post) { out.free(); return false; } out.ffn_residual = go.residual; @@ -146,8 +204,7 @@ static bool build_cached_attn_prefn( out.ctx, out.gf, w, cache, layer_idx, out.inp_embed, out.positions, /*attn_mask=*/nullptr, /*kv_start=*/kv_win - 1, /*n_tokens=*/1, /*fa_window=*/0, - out.kv_write_rows, - /*skip_gdn_intermediate=*/true); + out.kv_write_rows); if (!go.residual || !go.post) { out.free(); return false; } out.ffn_residual = go.residual; @@ -190,6 +247,15 @@ void PipelinedDecodeState::destroy() { routing_ids_buf.clear(); routing_weights_buf.clear(); ffn_post_host_buf.clear(); + hot_ids_buf.clear(); + hot_weights_buf.clear(); + cold_ids_buf.clear(); + cold_weights_buf.clear(); + moe_expert_layers.clear(); + moe_expert_output_buf.clear(); + moe_expert_runtime.reset(); + moe_multi_target_expert_runtime.reset(); + moe_expert_compute = nullptr; cold_in_zeroed = false; n_layer = 0; } @@ -197,9 +263,11 @@ void PipelinedDecodeState::destroy() { bool init_pipelined_decode_state( PipelinedDecodeState & out, ggml_backend_t backend, + const std::string & target_path, const TargetWeights & w, TargetCache & cache, MoeHybridStorage & hybrid, + const ExpertSplitComputeRuntime * split_runtime, int kv_start, int kq_stride_pad) { @@ -219,6 +287,10 @@ bool init_pipelined_decode_state( out.routing_ids_buf.resize((size_t)w.n_expert_used); out.routing_weights_buf.resize((size_t)w.n_expert_used); out.ffn_post_host_buf.resize((size_t)w.n_embd); + out.hot_ids_buf.resize((size_t)w.n_expert_used); + out.hot_weights_buf.resize((size_t)w.n_expert_used); + out.cold_ids_buf.resize((size_t)w.n_expert_used); + out.cold_weights_buf.resize((size_t)w.n_expert_used); // Check if routed FFN pipeline is disabled const bool routed_disabled = (std::getenv("DFLASH_QWEN35MOE_NO_ROUTED") != nullptr); @@ -231,7 +303,7 @@ bool init_pipelined_decode_state( out.cached_prefn.resize((size_t)w.n_layer); int cached_prefn_count = 0; for (int il = 0; il < w.n_layer; ++il) { - const bool is_attn = (((il + 1) % w.full_attention_interval) == 0); + const bool is_attn = qwen35_layer_is_full_attention(w, il); if (!is_attn) { if (!build_cached_deltanet_prefn( out.cached_prefn[(size_t)il], backend, w, cache, il, kv_start, kq_stride_pad)) { @@ -270,37 +342,53 @@ bool init_pipelined_decode_state( cached_prefn_count, routed_count, out.cold_compute ? "" : " (drop_cold=lossy)"); - // Initialize fused cold FFN compute (bypasses ggml graph dispatch) - if (out.cold_compute) { - out.cold_ffn_compute = make_cpu_cold_ffn_compute(w.n_ff_exp); - out.cold_ffn_layers.resize((size_t)w.n_layer); - out.cold_output_buf.resize((size_t)w.n_embd); - for (int il = 0; il < w.n_layer && (size_t)il < hybrid.layers.size(); ++il) { - auto & storage = hybrid.layers[(size_t)il]; - const TargetLayer & L = w.layers[(size_t)il]; - auto & cl = out.cold_ffn_layers[(size_t)il]; - cl.fused_gate_up = (storage.gate_up_cold != nullptr); - if (cl.fused_gate_up) { - cl.gate_up_data = storage.gate_up_cold ? storage.gate_up_cold->data : nullptr; - cl.gate_up_stride = storage.gate_up_cold ? storage.gate_up_cold->nb[2] : 0; - cl.gate_up_type = storage.gate_up_cold ? storage.gate_up_cold->type : GGML_TYPE_Q4_K; - cl.gate_up_scale = L.ffn_gate_up_exps_s; - } else { - cl.gate_data = storage.gate_cold ? storage.gate_cold->data : nullptr; - cl.up_data = storage.up_cold ? storage.up_cold->data : nullptr; - cl.gate_stride = storage.gate_cold ? storage.gate_cold->nb[2] : 0; - cl.up_stride = storage.up_cold ? storage.up_cold->nb[2] : 0; - cl.gate_type = storage.gate_cold ? storage.gate_cold->type : GGML_TYPE_Q4_K; - cl.up_type = storage.up_cold ? storage.up_cold->type : GGML_TYPE_Q4_K; - cl.gate_scale = L.ffn_gate_exps_s; - cl.up_scale = L.ffn_up_exps_s; - } - cl.down_data = storage.down_cold ? storage.down_cold->data : nullptr; - cl.down_stride = storage.down_cold ? storage.down_cold->nb[2] : 0; - cl.down_type = storage.down_cold ? storage.down_cold->type : GGML_TYPE_Q4_K; - cl.down_scale = L.ffn_down_exps_s; + MoeExpertComputeRuntimeConfig runtime_cfg; + runtime_cfg.target_path = target_path; + runtime_cfg.n_layer = w.n_layer; + runtime_cfg.n_expert = w.n_expert; + runtime_cfg.n_expert_used = w.n_expert_used; + runtime_cfg.n_embd = w.n_embd; + runtime_cfg.n_ff_exp = w.n_ff_exp; + runtime_cfg.enabled = out.cold_compute; + runtime_cfg.log_prefix = "[pipelined]"; + + std::vector layer_descs((size_t)w.n_layer); + for (int il = 0; il < w.n_layer; ++il) { + layer_descs[(size_t)il] = make_moe_layer_desc(w.layers[(size_t)il]); + } + + std::string expert_err; + bool multi_target_enabled = false; + if (split_runtime && split_runtime->matches(w.n_layer, w.n_expert, w.n_expert_used)) { + if (!ensure_multi_target_moe_expert_compute_runtime( + out.moe_multi_target_expert_runtime, runtime_cfg, *split_runtime, + hybrid, layer_descs, &expert_err)) { + std::fprintf(stderr, "[pipelined] multi-target MoE expert runtime init failed: %s\n", + expert_err.c_str()); + return false; } - std::fprintf(stderr, "[pipelined] cold FFN: fused kernel (bypasses ggml graph)\n"); + multi_target_enabled = out.moe_multi_target_expert_runtime.enabled && + out.moe_multi_target_expert_runtime.compute_ptr() != nullptr; + } + if (!multi_target_enabled && + !ensure_moe_expert_compute_runtime(out.moe_expert_runtime, runtime_cfg, + hybrid, layer_descs, &expert_err)) { + std::fprintf(stderr, "[pipelined] MoE expert runtime init failed: %s\n", + expert_err.c_str()); + return false; + } + out.moe_expert_compute = multi_target_enabled + ? out.moe_multi_target_expert_runtime.compute_ptr() + : out.moe_expert_runtime.compute_ptr(); + out.moe_expert_layers = multi_target_enabled + ? out.moe_multi_target_expert_runtime.layers + : out.moe_expert_runtime.layers; + out.moe_expert_output_buf.resize((size_t)w.n_embd); + if (out.moe_expert_compute) { + std::fprintf(stderr, "[pipelined] MoE expert compute ready: %s\n", + multi_target_enabled + ? out.moe_multi_target_expert_runtime.runtime_key.c_str() + : out.moe_expert_runtime.runtime_key.c_str()); } out.cold_in_zeroed = true; @@ -332,11 +420,23 @@ bool pipelined_decode_one_token( *tel = PipelinedDecodeTelemetry{}; } + if (state.routing_ids_buf.size() < (size_t)n_expert_used || + state.routing_weights_buf.size() < (size_t)n_expert_used || + state.hot_ids_buf.size() < (size_t)n_expert_used || + state.hot_weights_buf.size() < (size_t)n_expert_used || + state.cold_ids_buf.size() < (size_t)n_expert_used || + state.cold_weights_buf.size() < (size_t)n_expert_used) { + std::fprintf(stderr, + "[pipelined] routing scratch under-sized for n_expert_used=%d\n", + n_expert_used); + return false; + } + const auto tok_t0 = PipelineClock::now(); StepGraph dyn_sg; // for attention layers (rebuilt per-token) for (int il = 0; il < n_layer; ++il) { - const bool is_attn = (((il + 1) % state.full_attention_interval) == 0); + const bool is_attn = qwen35_layer_is_full_attention(w, il); const auto prefn_build_t0 = PipelineClock::now(); // ══════════════════════════════════════════════════════════════════════ @@ -367,11 +467,11 @@ bool pipelined_decode_one_token( const auto sync_t1 = PipelineClock::now(); // Read routing decisions from GPU - int32_t global_ids[8]; - float router_weights[8]; - ggml_backend_tensor_get(cpg.moe_selected, global_ids, 0, + auto & global_ids = state.routing_ids_buf; + auto & router_weights = state.routing_weights_buf; + ggml_backend_tensor_get(cpg.moe_selected, global_ids.data(), 0, sizeof(int32_t) * (size_t)n_expert_used); - ggml_backend_tensor_get(cpg.moe_weights, router_weights, 0, + ggml_backend_tensor_get(cpg.moe_weights, router_weights.data(), 0, sizeof(float) * (size_t)n_expert_used); const auto readback_t1 = PipelineClock::now(); @@ -382,44 +482,32 @@ bool pipelined_decode_one_token( sizeof(float) * (size_t)n_embd); ffn_post_on_host = true; state.routing_collector->record(il, state.ffn_post_host_buf.data(), - n_embd, global_ids, n_expert_used); + n_embd, global_ids.data(), n_expert_used); } // CPU-side local ID mapping + cold partition (trivial: 8 lookups) auto & storage = hybrid.layers[(size_t)il]; - int32_t local_ids[8]; - float masked_weights[8]; - int32_t cold_ids[8]; - float cold_weights[8]; + auto & local_ids = state.hot_ids_buf; + auto & masked_weights = state.hot_weights_buf; + auto & cold_ids = state.cold_ids_buf; + auto & cold_weights = state.cold_weights_buf; int n_cold = 0; int layer_cold_hits = 0; // Spark expert cache: pull selected cold experts into spare GPU slots // (LRU) so the lookup below serves them on-GPU; after warmup cold->0. if (storage.cache_slots > 0) for (int i = 0; i < n_expert_used; ++i) - dflash::common::moe_hybrid_cache_swap_in(storage, global_ids[i], backend); - for (int i = 0; i < n_expert_used; ++i) { - int32_t gid = global_ids[i]; - int32_t lid = (gid >= 0 && gid < (int)storage.hot_local_by_global.size()) - ? storage.hot_local_by_global[(size_t)gid] : -1; - if (lid >= 0) { - local_ids[i] = lid; - masked_weights[i] = router_weights[i]; - } else { - local_ids[i] = 0; // safe: maps to expert 0 (result zeroed by weight) - masked_weights[i] = 0.0f; // cold expert contributes nothing to hot path - layer_cold_hits++; - // Record for cold compute - if (state.cold_ffn_compute && gid >= 0 && gid < (int)storage.cold_local_by_global.size()) { - int32_t cold_local = storage.cold_local_by_global[(size_t)gid]; - if (cold_local >= 0) { - cold_ids[n_cold] = cold_local; - cold_weights[n_cold] = router_weights[i]; - n_cold++; - } - } - } + dflash::common::moe_hybrid_cache_swap_in(storage, global_ids[(size_t)i], backend); + if (!prepare_unique_decode_hot_slots( + storage, global_ids.data(), router_weights.data(), + n_expert_used, local_ids, masked_weights, + cold_ids, cold_weights, n_cold)) { + std::fprintf(stderr, + "[pipelined] invalid hot/cold routed selection in layer %d\n", + il); + return false; } + layer_cold_hits = n_cold; const bool has_cold_selected = (n_cold > 0); const auto remap_t1 = PipelineClock::now(); @@ -432,9 +520,9 @@ bool pipelined_decode_one_token( } // Upload pre-computed inputs to rffn graph (H→D async on compute stream) - ggml_backend_tensor_set_async(backend, rffn.ids, local_ids, 0, + ggml_backend_tensor_set_async(backend, rffn.ids, local_ids.data(), 0, sizeof(int32_t) * (size_t)n_expert_used); - ggml_backend_tensor_set_async(backend, rffn.weights, masked_weights, 0, + ggml_backend_tensor_set_async(backend, rffn.weights, masked_weights.data(), 0, sizeof(float) * (size_t)n_expert_used); // Copy ffn_post from prefn output → rffn input (GPU→GPU, already synced) ggml_backend_tensor_copy_async(backend, backend, cpg.ffn_post, rffn.inp); @@ -448,12 +536,14 @@ bool pipelined_decode_one_token( // 6. Cold compute on CPU (parallel with GPU rffn above) const auto cold_t0 = PipelineClock::now(); if (has_cold_selected) { - state.cold_ffn_compute->compute( - state.cold_ffn_layers[(size_t)il], + if (!state.moe_expert_compute->compute( + state.moe_expert_layers[(size_t)il], state.ffn_post_host_buf.data(), - cold_ids, cold_weights, n_cold, + cold_ids.data(), cold_weights.data(), n_cold, n_embd, w.n_ff_exp, - state.cold_output_buf.data()); + state.moe_expert_output_buf.data())) { + return false; + } } if (tel && has_cold_selected) tel->cold_compute_us += pipe_elapsed_us(cold_t0, PipelineClock::now()); @@ -463,7 +553,7 @@ bool pipelined_decode_one_token( // 8. Upload cold result or ensure cold_in is zero if (has_cold_selected) { ggml_backend_tensor_set_async(backend, state.gpu_state.combine.cold_in, - state.cold_output_buf.data(), 0, + state.moe_expert_output_buf.data(), 0, sizeof(float) * (size_t)n_embd); state.cold_in_zeroed = false; } else if (!state.cold_in_zeroed) { @@ -649,37 +739,26 @@ bool pipelined_decode_one_token( auto & rffn = state.cached_routed_ffn[(size_t)il]; if (rffn.valid()) { // Partition hot/cold: remap global→local, zero cold weights for hot path - int32_t local_ids[8]; - float masked_weights[8]; - int32_t cold_ids[8]; - float cold_weights[8]; + auto & local_ids = state.hot_ids_buf; + auto & masked_weights = state.hot_weights_buf; + auto & cold_ids = state.cold_ids_buf; + auto & cold_weights = state.cold_weights_buf; int n_cold = 0; int layer_cold_hits = 0; // Spark expert cache: pull selected cold experts into spare GPU slots. if (storage.cache_slots > 0) for (int i = 0; i < n_expert_used; ++i) dflash::common::moe_hybrid_cache_swap_in(storage, state.routing_ids_buf[(size_t)i], backend); - for (int i = 0; i < n_expert_used; ++i) { - int32_t gid = state.routing_ids_buf[(size_t)i]; - int32_t lid = (gid >= 0 && gid < (int)storage.hot_local_by_global.size()) - ? storage.hot_local_by_global[(size_t)gid] : -1; - if (lid >= 0) { - local_ids[i] = lid; - masked_weights[i] = state.routing_weights_buf[(size_t)i]; - } else { - local_ids[i] = 0; - masked_weights[i] = 0.0f; - layer_cold_hits++; - if (state.cold_ffn_compute && gid >= 0 && gid < (int)storage.cold_local_by_global.size()) { - int32_t cold_local = storage.cold_local_by_global[(size_t)gid]; - if (cold_local >= 0) { - cold_ids[n_cold] = cold_local; - cold_weights[n_cold] = state.routing_weights_buf[(size_t)i]; - n_cold++; - } - } - } + if (!prepare_unique_decode_hot_slots( + storage, state.routing_ids_buf.data(), state.routing_weights_buf.data(), + n_expert_used, local_ids, masked_weights, + cold_ids, cold_weights, n_cold)) { + std::fprintf(stderr, + "[pipelined] invalid hot/cold routed selection in layer %d\n", + il); + return false; } + layer_cold_hits = n_cold; const bool has_cold_selected = (n_cold > 0); // D2H ffn_post for cold compute (GPU already synced after routing readback). @@ -690,9 +769,9 @@ bool pipelined_decode_one_token( } // Upload IDs + weights, copy inputs, dispatch rffn (all async) - ggml_backend_tensor_set_async(backend, rffn.ids, local_ids, 0, + ggml_backend_tensor_set_async(backend, rffn.ids, local_ids.data(), 0, sizeof(int32_t) * (size_t)n_expert_used); - ggml_backend_tensor_set_async(backend, rffn.weights, masked_weights, 0, + ggml_backend_tensor_set_async(backend, rffn.weights, masked_weights.data(), 0, sizeof(float) * (size_t)n_expert_used); ggml_backend_tensor_copy_async(backend, backend, ffn_post_gpu, rffn.inp); ggml_backend_tensor_copy_async(backend, backend, ffn_residual_gpu, state.gpu_state.combine.residual_in); @@ -701,12 +780,14 @@ bool pipelined_decode_one_token( // Cold compute on CPU (parallel with GPU rffn above) const auto cold_t0 = PipelineClock::now(); if (has_cold_selected) { - state.cold_ffn_compute->compute( - state.cold_ffn_layers[(size_t)il], + if (!state.moe_expert_compute->compute( + state.moe_expert_layers[(size_t)il], state.ffn_post_host_buf.data(), - cold_ids, cold_weights, n_cold, + cold_ids.data(), cold_weights.data(), n_cold, n_embd, w.n_ff_exp, - state.cold_output_buf.data()); + state.moe_expert_output_buf.data())) { + return false; + } } if (tel && has_cold_selected) tel->cold_compute_us += pipe_elapsed_us(cold_t0, PipelineClock::now()); @@ -716,7 +797,7 @@ bool pipelined_decode_one_token( // Upload cold result or ensure cold_in is zero if (has_cold_selected) { ggml_backend_tensor_set_async(backend, state.gpu_state.combine.cold_in, - state.cold_output_buf.data(), 0, + state.moe_expert_output_buf.data(), 0, sizeof(float) * (size_t)n_embd); state.cold_in_zeroed = false; } else if (!state.cold_in_zeroed) { @@ -751,22 +832,24 @@ bool pipelined_decode_one_token( // Partition into hot/cold (fast: just a lookup table scan, ~8 iterations) int n_hot = 0, n_cold = 0; - int32_t hot_ids[8], cold_ids[8]; - float hot_weights[8], cold_weights[8]; + auto & hot_ids = state.hot_ids_buf; + auto & hot_weights = state.hot_weights_buf; + auto & cold_ids = state.cold_ids_buf; + auto & cold_weights = state.cold_weights_buf; for (int i = 0; i < n_expert_used; ++i) { const int32_t gid = state.routing_ids_buf[(size_t)i]; if (gid < 0 || gid >= (int32_t)storage.hot_local_by_global.size()) return false; const int32_t hot_local = storage.hot_local_by_global[(size_t)gid]; if (hot_local >= 0) { - hot_ids[n_hot] = hot_local; - hot_weights[n_hot] = state.routing_weights_buf[(size_t)i]; + hot_ids[(size_t)n_hot] = hot_local; + hot_weights[(size_t)n_hot] = state.routing_weights_buf[(size_t)i]; n_hot++; } else { const int32_t cold_local = storage.cold_local_by_global[(size_t)gid]; if (cold_local >= 0) { - cold_ids[n_cold] = cold_local; - cold_weights[n_cold] = state.routing_weights_buf[(size_t)i]; + cold_ids[(size_t)n_cold] = cold_local; + cold_weights[(size_t)n_cold] = state.routing_weights_buf[(size_t)i]; n_cold++; } } @@ -806,11 +889,11 @@ bool pipelined_decode_one_token( // All setup on compute stream — no per-op cudaStreamSynchronize ggml_backend_tensor_copy_async(backend, backend, ffn_post_gpu, storage.hot_graph.inp); if (storage.hot_graph.ids && has_hot) { - ggml_backend_tensor_set_async(backend, storage.hot_graph.ids, hot_ids, 0, + ggml_backend_tensor_set_async(backend, storage.hot_graph.ids, hot_ids.data(), 0, sizeof(int32_t) * (size_t)n_hot); } if (storage.hot_graph.weights && has_hot) { - ggml_backend_tensor_set_async(backend, storage.hot_graph.weights, hot_weights, 0, + ggml_backend_tensor_set_async(backend, storage.hot_graph.weights, hot_weights.data(), 0, sizeof(float) * (size_t)n_hot); } // Launch hot GPU async — queued after copies on same stream @@ -825,15 +908,18 @@ bool pipelined_decode_one_token( if (has_cold) { // ffn_post already read above (before hot launch) — no GPU sync here! const auto cold_compute_t0 = PipelineClock::now(); - if (state.cold_ffn_compute) { - // Fused kernel: bypass ggml graph dispatch entirely - state.cold_ffn_compute->compute( - state.cold_ffn_layers[(size_t)il], + if (state.moe_expert_compute) { + // Fused compute: bypass ggml graph dispatch entirely. + if (!state.moe_expert_compute->compute( + state.moe_expert_layers[(size_t)il], state.ffn_post_host_buf.data(), - cold_ids, - cold_weights, + cold_ids.data(), + cold_weights.data(), n_cold, n_embd, w.n_ff_exp, - state.cold_output_buf.data()); + state.moe_expert_output_buf.data())) { + if (hot_async_launched) ggml_backend_synchronize(backend); + return false; + } } else { // Fallback: ggml cold graph (legacy path) if (!storage.cold_graph.valid() || storage.cold_graph.n_hot != n_cold) { @@ -845,9 +931,9 @@ bool pipelined_decode_one_token( if (storage.cold_graph.valid() && storage.cold_graph.n_hot == n_cold) { ggml_backend_tensor_set(storage.cold_graph.inp, state.ffn_post_host_buf.data(), 0, sizeof(float) * (size_t)n_embd); - ggml_backend_tensor_set(storage.cold_graph.ids, cold_ids, 0, + ggml_backend_tensor_set(storage.cold_graph.ids, cold_ids.data(), 0, sizeof(int32_t) * (size_t)n_cold); - ggml_backend_tensor_set(storage.cold_graph.weights, cold_weights, 0, + ggml_backend_tensor_set(storage.cold_graph.weights, cold_weights.data(), 0, sizeof(float) * (size_t)n_cold); auto cst = ggml_backend_graph_compute(cpu_be, storage.cold_graph.gf); if (cst != GGML_STATUS_SUCCESS) { @@ -875,8 +961,8 @@ bool pipelined_decode_one_token( } if (has_cold) { - const float * cold_result = state.cold_ffn_compute - ? state.cold_output_buf.data() + const float * cold_result = state.moe_expert_compute + ? state.moe_expert_output_buf.data() : nullptr; if (!cold_result) { // Legacy path: read from ggml tensor diff --git a/server/src/qwen35moe/qwen35moe_pipelined_decode.h b/server/src/qwen35moe/qwen35moe_pipelined_decode.h index af342350f..852a8a807 100644 --- a/server/src/qwen35moe/qwen35moe_pipelined_decode.h +++ b/server/src/qwen35moe/qwen35moe_pipelined_decode.h @@ -13,7 +13,7 @@ #include "../common/moe_hybrid_ffn_eval.h" #include "../common/moe_hybrid_storage.h" #include "../common/moe_routing_collector.h" -#include "../common/cold_ffn_compute.h" +#include "../common/moe_expert_compute.h" #include "graph_builders.h" #include "ggml-backend.h" @@ -118,6 +118,10 @@ struct PipelinedDecodeState { std::vector routing_ids_buf; std::vector routing_weights_buf; std::vector ffn_post_host_buf; + std::vector hot_ids_buf; + std::vector hot_weights_buf; + std::vector cold_ids_buf; + std::vector cold_weights_buf; // Persistent zero buffer for cold_in (set once at init) bool cold_in_zeroed = false; @@ -127,10 +131,14 @@ struct PipelinedDecodeState { // Set DFLASH_DROP_COLD=1 to disable (fast but lossy). bool cold_compute = true; - // Fused cold FFN compute (bypasses ggml graph dispatch overhead) - std::unique_ptr cold_ffn_compute; - std::vector cold_ffn_layers; // per-layer cold weight metadata - std::vector cold_output_buf; // [n_embd] scratch for cold FFN output + // Shared expert runtime: CPU fallback or remote IPC daemon. + MoeExpertComputeRuntime moe_expert_runtime; + MoeMultiTargetExpertRuntime moe_multi_target_expert_runtime; + + // Fused selected-expert compute (non-owning view into the active runtime). + MoeExpertCompute * moe_expert_compute = nullptr; + std::vector moe_expert_layers; // per-layer expert weight metadata + std::vector moe_expert_output_buf; // [n_embd] scratch for expert output // Tracking int n_layer = 0; @@ -152,11 +160,20 @@ struct PipelinedDecodeState { routing_ids_buf(std::move(o.routing_ids_buf)), routing_weights_buf(std::move(o.routing_weights_buf)), ffn_post_host_buf(std::move(o.ffn_post_host_buf)), + hot_ids_buf(std::move(o.hot_ids_buf)), + hot_weights_buf(std::move(o.hot_weights_buf)), + cold_ids_buf(std::move(o.cold_ids_buf)), + cold_weights_buf(std::move(o.cold_weights_buf)), cold_in_zeroed(o.cold_in_zeroed), cold_compute(o.cold_compute), - cold_ffn_compute(std::move(o.cold_ffn_compute)), - cold_ffn_layers(std::move(o.cold_ffn_layers)), - cold_output_buf(std::move(o.cold_output_buf)), + moe_expert_runtime(std::move(o.moe_expert_runtime)), + moe_multi_target_expert_runtime(std::move(o.moe_multi_target_expert_runtime)), + moe_expert_compute( + moe_multi_target_expert_runtime.compute_ptr() + ? moe_multi_target_expert_runtime.compute_ptr() + : moe_expert_runtime.compute_ptr()), + moe_expert_layers(std::move(o.moe_expert_layers)), + moe_expert_output_buf(std::move(o.moe_expert_output_buf)), n_layer(o.n_layer), n_embd(o.n_embd), n_expert_used(o.n_expert_used), full_attention_interval(o.full_attention_interval) { @@ -171,11 +188,21 @@ struct PipelinedDecodeState { routing_ids_buf = std::move(o.routing_ids_buf); routing_weights_buf = std::move(o.routing_weights_buf); ffn_post_host_buf = std::move(o.ffn_post_host_buf); + hot_ids_buf = std::move(o.hot_ids_buf); + hot_weights_buf = std::move(o.hot_weights_buf); + cold_ids_buf = std::move(o.cold_ids_buf); + cold_weights_buf = std::move(o.cold_weights_buf); cold_in_zeroed = o.cold_in_zeroed; cold_compute = o.cold_compute; - cold_ffn_compute = std::move(o.cold_ffn_compute); - cold_ffn_layers = std::move(o.cold_ffn_layers); - cold_output_buf = std::move(o.cold_output_buf); + moe_expert_runtime = std::move(o.moe_expert_runtime); + moe_multi_target_expert_runtime = + std::move(o.moe_multi_target_expert_runtime); + moe_expert_compute = + moe_multi_target_expert_runtime.compute_ptr() + ? moe_multi_target_expert_runtime.compute_ptr() + : moe_expert_runtime.compute_ptr(); + moe_expert_layers = std::move(o.moe_expert_layers); + moe_expert_output_buf = std::move(o.moe_expert_output_buf); n_layer = o.n_layer; n_embd = o.n_embd; n_expert_used = o.n_expert_used; full_attention_interval = o.full_attention_interval; @@ -192,9 +219,11 @@ struct PipelinedDecodeState { bool init_pipelined_decode_state( PipelinedDecodeState & out, ggml_backend_t backend, + const std::string & target_path, const TargetWeights & w, TargetCache & cache, MoeHybridStorage & hybrid, + const ExpertSplitComputeRuntime * split_runtime, int kv_start, // initial KV position for graph caching int kq_stride_pad); diff --git a/server/test/test_expert_split_materialization.cpp b/server/test/test_expert_split_materialization.cpp new file mode 100644 index 000000000..7823aaa98 --- /dev/null +++ b/server/test/test_expert_split_materialization.cpp @@ -0,0 +1,104 @@ +#include "../src/common/expert_split_materialization.h" +#include "../src/common/expert_split_plan.h" +#include "../src/common/expert_split_runtime.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +static void expect(bool cond, const char * msg) { + if (!cond) { + std::fprintf(stderr, "FAIL: %s\n", msg); + std::exit(1); + } +} + +static ExpertSplitPlan build_multi_target_plan() { + ExpertSplitConfig cfg; + cfg.n_layer = 2; + cfg.n_expert = 4; + + const std::vector targets = { + {"cuda:0", "cuda", 0, 4, 0, false}, + {"hip:0", "hip", 0, 2, 0, false}, + }; + + const std::vector units = { + {0, 0, 1, 100.0}, {0, 1, 1, 90.0}, {0, 2, 1, 80.0}, {0, 3, 1, 70.0}, + {1, 0, 1, 60.0}, {1, 1, 1, 50.0}, {1, 2, 1, 40.0}, {1, 3, 1, 30.0}, + }; + + ExpertSplitPlan plan; + std::string err; + expect(build_expert_split_plan(cfg, targets, units, plan, &err), err.c_str()); + return plan; +} + +static void test_materialization_builds_ordered_cold_union() { + const ExpertSplitPlan plan = build_multi_target_plan(); + ExpertSplitRuntime runtime; + std::string err; + expect(build_expert_split_runtime(plan, runtime, &err), err.c_str()); + + ExpertSplitMaterialization materialization; + expect(build_expert_split_materialization(runtime, /*n_expert_used=*/2, + materialization, &err), + err.c_str()); + + expect(materialization.matches(2, 4, 2), "materialization dims"); + expect(materialization.ordered_cold_union, "ordered union enabled"); + expect(materialization.targets.size() == 3, "targets include cpu fallback"); + expect(materialization.primary_placement.total_hot == 4, "primary hot total"); + + const auto & l0_hot = materialization.primary_placement.hot_expert_ids[0]; + const auto & l1_hot = materialization.primary_placement.hot_expert_ids[1]; + expect(l0_hot.size() == 4, "layer0 primary hot count"); + expect(l1_hot.empty(), "layer1 primary hot empty"); + expect(l0_hot[0] == 0 && l0_hot[1] == 1 && l0_hot[2] == 2 && l0_hot[3] == 3, + "layer0 primary hot order"); + + const auto & l0_cold = materialization.cold_expert_ids_by_layer[0]; + const auto & l1_cold = materialization.cold_expert_ids_by_layer[1]; + expect(l0_cold.empty(), "layer0 cold empty"); + expect(l1_cold.size() == 4, "layer1 cold count"); + expect(l1_cold[0] == 0 && l1_cold[1] == 1 && l1_cold[2] == 2 && + l1_cold[3] == 3, "layer1 cold preserves target order"); +} + +static void test_materialization_single_explicit_gpu_keeps_single_primary_shape() { + ExpertSplitConfig cfg; + cfg.n_layer = 1; + cfg.n_expert = 4; + + const std::vector targets = { + {"cuda:0", "cuda", 0, 2, 0, false}, + }; + const std::vector units = { + {0, 0, 1, 40.0}, {0, 1, 1, 30.0}, {0, 2, 1, 20.0}, {0, 3, 1, 10.0}, + }; + + ExpertSplitPlan plan; + std::string err; + expect(build_expert_split_plan(cfg, targets, units, plan, &err), err.c_str()); + + ExpertSplitRuntime runtime; + expect(build_expert_split_runtime(plan, runtime, &err), err.c_str()); + + ExpertSplitMaterialization materialization; + expect(build_expert_split_materialization(runtime, /*n_expert_used=*/2, + materialization, &err), + err.c_str()); + expect(materialization.matches(1, 4, 2), "single-primary materialization dims"); + expect(!materialization.ordered_cold_union, "single-primary path keeps existing cold order"); + expect(materialization.cold_expert_ids_by_layer.empty(), "no cold override"); +} + +int main() { + test_materialization_builds_ordered_cold_union(); + test_materialization_single_explicit_gpu_keeps_single_primary_shape(); + std::printf("OK\n"); + return 0; +} diff --git a/server/test/test_expert_split_plan.cpp b/server/test/test_expert_split_plan.cpp new file mode 100644 index 000000000..99e9933bc --- /dev/null +++ b/server/test/test_expert_split_plan.cpp @@ -0,0 +1,81 @@ +#include "../src/common/expert_split_plan.h" +#include "../src/common/moe_hybrid_placement.h" +#include "../src/common/moe_hybrid_routing_stats.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +static void expect(bool cond, const char * msg) { + if (!cond) { + std::fprintf(stderr, "FAIL: %s\n", msg); + std::exit(1); + } +} + +static void test_ordered_targets_and_cpu_fallback() { + ExpertSplitConfig cfg; + cfg.n_layer = 1; + cfg.n_expert = 5; + + const std::vector targets = { + {"cuda:0", "cuda", 0, 2, 0, false}, + {"hip:0", "hip", 0, 1, 0, false}, + }; + const std::vector units = { + {0, 0, 1, 50.0}, + {0, 1, 1, 40.0}, + {0, 2, 1, 30.0}, + {0, 3, 1, 20.0}, + {0, 4, 1, 10.0}, + }; + + ExpertSplitPlan plan; + std::string err; + expect(build_expert_split_plan(cfg, targets, units, plan, &err), err.c_str()); + expect(plan.matches(1, 5), "plan matches dims"); + expect((int) plan.targets.size() == 3, "implicit cpu fallback added"); + expect(plan.at(0, 0).target_index == 0, "best expert on first target"); + expect(plan.at(0, 1).target_index == 0, "second expert on first target"); + expect(plan.at(0, 2).target_index == 1, "third expert on second target"); + expect(plan.at(0, 3).target_index == 2, "overflow expert on cpu"); + expect(plan.at(0, 4).target_index == 2, "tail expert on cpu"); +} + +static void test_hybrid_placement_uses_expert_split_planner() { + MoeHybridRoutingStats stats; + stats.n_layer = 2; + stats.n_expert = 4; + stats.n_expert_used = 2; + stats.counts = { + 100, 80, 60, 40, + 100, 1, 1, 1, + }; + stats.layer_totals = {280, 103}; + + const std::vector layer_bytes = {1, 1}; + MoeHybridPlacement placement; + std::string err; + expect(MoeHybridPlacement::build_from_stats_with_layer_bytes( + stats, layer_bytes, /*total_hot_budget_bytes=*/4, + /*min_hot_per_layer=*/1, placement, &err), + err.c_str()); + expect(placement.hot_counts.size() == 2, "hot counts size"); + expect(placement.hot_counts[0] == 3, "layer0 got extra hot slots"); + expect(placement.hot_counts[1] == 1, "layer1 kept minimum hot slot"); + expect(placement.is_hot(0, 0), "layer0 expert0 hot"); + expect(placement.is_hot(0, 1), "layer0 expert1 hot"); + expect(placement.is_hot(0, 2), "layer0 expert2 hot"); + expect(!placement.is_hot(0, 3), "layer0 expert3 cold"); + expect(placement.is_hot(1, 0), "layer1 expert0 hot"); +} + +int main() { + test_ordered_targets_and_cpu_fallback(); + test_hybrid_placement_uses_expert_split_planner(); + std::printf("OK\n"); + return 0; +} diff --git a/server/test/test_expert_split_runtime.cpp b/server/test/test_expert_split_runtime.cpp new file mode 100644 index 000000000..cfa7c6213 --- /dev/null +++ b/server/test/test_expert_split_runtime.cpp @@ -0,0 +1,257 @@ +#include "../src/common/expert_split_plan.h" +#include "../src/common/expert_split_compute_runtime.h" +#include "../src/common/expert_split_runtime.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +static void expect(bool cond, const char * msg) { + if (!cond) { + std::fprintf(stderr, "FAIL: %s\n", msg); + std::exit(1); + } +} + +static ExpertSplitPlan build_sample_plan() { + ExpertSplitConfig cfg; + cfg.n_layer = 2; + cfg.n_expert = 4; + + const std::vector targets = { + {"cuda:0", "cuda", 0, 4, 0, false}, + {"hip:0", "hip", 0, 2, 0, false}, + }; + const std::vector units = { + {0, 0, 1, 100.0}, + {0, 1, 1, 90.0}, + {0, 2, 1, 80.0}, + {0, 3, 1, 70.0}, + {1, 0, 1, 60.0}, + {1, 1, 1, 50.0}, + {1, 2, 1, 40.0}, + {1, 3, 1, 30.0}, + }; + + ExpertSplitPlan plan; + std::string err; + expect(build_expert_split_plan(cfg, targets, units, plan, &err), err.c_str()); + return plan; +} + +static void test_runtime_materializes_targets_and_cpu_fallback() { + const ExpertSplitPlan plan = build_sample_plan(); + ExpertSplitRuntime runtime; + std::string err; + expect(build_expert_split_runtime(plan, runtime, &err), err.c_str()); + + expect(runtime.matches(2, 4), "runtime dims"); + expect((int) runtime.targets.size() == 3, "runtime keeps cpu fallback target"); + expect(runtime.targets[0].target.name == "cuda:0", "target0 name"); + expect(runtime.targets[1].target.name == "hip:0", "target1 name"); + expect(runtime.targets[2].target.backend == "cpu", "target2 cpu fallback"); + + expect(runtime.target_index(0, 0) == 0, "layer0 expert0 on target0"); + expect(runtime.target_index(0, 1) == 0, "layer0 expert1 on target0"); + expect(runtime.target_index(0, 2) == 0, "layer0 expert2 on target0"); + expect(runtime.target_index(0, 3) == 0, "layer0 expert3 on target0"); + expect(runtime.target_index(1, 0) == 1, "layer1 expert0 on target1"); + expect(runtime.target_index(1, 1) == 1, "layer1 expert1 on target1"); + expect(runtime.target_index(1, 2) == 2, "layer1 expert2 on cpu"); + expect(runtime.target_index(1, 3) == 2, "layer1 expert3 on cpu"); + + expect(runtime.targets[0].used_bytes == 4, "target0 bytes"); + expect(runtime.targets[1].used_bytes == 2, "target1 bytes"); + expect(runtime.targets[2].used_bytes == 2, "target2 bytes"); +} + +static void test_runtime_local_order_is_score_stable() { + const ExpertSplitPlan plan = build_sample_plan(); + ExpertSplitRuntime runtime; + std::string err; + expect(build_expert_split_runtime(plan, runtime, &err), err.c_str()); + + const ExpertSplitLayerTarget * layer0_t0 = runtime.layer_target_ptr(0, 0); + expect(layer0_t0 != nullptr, "layer0 target0 exists"); + expect(layer0_t0->global_expert_ids.size() == 4, "layer0 target0 expert count"); + expect(layer0_t0->global_expert_ids[0] == 0, "layer0 target0 local0"); + expect(layer0_t0->global_expert_ids[1] == 1, "layer0 target0 local1"); + expect(layer0_t0->global_expert_ids[2] == 2, "layer0 target0 local2"); + expect(layer0_t0->global_expert_ids[3] == 3, "layer0 target0 local3"); + expect(runtime.local_index(0, 0) == 0, "layer0 expert0 local"); + expect(runtime.local_index(0, 3) == 3, "layer0 expert3 local"); + + const ExpertSplitLayerTarget * layer1_t2 = runtime.layer_target_ptr(2, 1); + expect(layer1_t2 != nullptr, "layer1 cpu exists"); + expect(layer1_t2->global_expert_ids.size() == 2, "layer1 cpu expert count"); + expect(layer1_t2->global_expert_ids[0] == 2, "layer1 cpu local0 by score"); + expect(layer1_t2->global_expert_ids[1] == 3, "layer1 cpu local1 by score"); + expect(runtime.local_index(1, 2) == 0, "layer1 expert2 local"); + expect(runtime.local_index(1, 3) == 1, "layer1 expert3 local"); +} + +static void test_target_placement_roundtrip() { + const ExpertSplitPlan plan = build_sample_plan(); + ExpertSplitRuntime runtime; + std::string err; + expect(build_expert_split_runtime(plan, runtime, &err), err.c_str()); + + MoeHybridPlacement t0; + expect(build_expert_split_target_placement(runtime, /*n_expert_used=*/2, 0, t0, &err), + err.c_str()); + expect(t0.matches(2, 4, 2), "target0 placement dims"); + expect(t0.total_hot == 4, "target0 total hot"); + expect(t0.hot_counts[0] == 4, "target0 layer0 count"); + expect(t0.hot_counts[1] == 0, "target0 layer1 count"); + expect(t0.hot_expert_ids[0].size() == 4, "target0 layer0 ids"); + expect(t0.hot_expert_ids[1].empty(), "target0 layer1 empty"); + + MoeHybridPlacement t2; + expect(build_expert_split_target_placement(runtime, /*n_expert_used=*/2, 2, t2, &err), + err.c_str()); + expect(t2.total_hot == 2, "target2 total hot"); + expect(t2.hot_counts[0] == 0, "target2 layer0 count"); + expect(t2.hot_counts[1] == 2, "target2 layer1 count"); + expect(t2.hot_expert_ids[1].size() == 2, "target2 layer1 ids"); + expect(t2.hot_expert_ids[1][0] == 2, "target2 layer1 first"); + expect(t2.hot_expert_ids[1][1] == 3, "target2 layer1 second"); +} + +static void test_all_target_placements_materialize() { + const ExpertSplitPlan plan = build_sample_plan(); + ExpertSplitRuntime runtime; + std::string err; + expect(build_expert_split_runtime(plan, runtime, &err), err.c_str()); + + std::vector placements; + expect(build_all_expert_split_target_placements(runtime, /*n_expert_used=*/2, + placements, &err), + err.c_str()); + expect(placements.size() == 3, "all target placements count"); + expect(placements[0].target_index == 0, "placement0 index"); + expect(placements[1].target_index == 1, "placement1 index"); + expect(placements[2].target_index == 2, "placement2 index"); + expect(placements[0].placement.total_hot == 4, "placement0 total hot"); + expect(placements[1].placement.total_hot == 2, "placement1 total hot"); + expect(placements[2].placement.total_hot == 2, "placement2 total hot"); +} + +static void test_compute_runtime_materializes_target_views() { + const ExpertSplitPlan plan = build_sample_plan(); + ExpertSplitRuntime runtime; + std::string err; + expect(build_expert_split_runtime(plan, runtime, &err), err.c_str()); + + ExpertSplitComputeRuntime compute_runtime; + expect(build_expert_split_compute_runtime(runtime, /*n_expert_used=*/2, + compute_runtime, &err), + err.c_str()); + expect(compute_runtime.matches(2, 4, 2), "compute runtime dims"); + expect(compute_runtime.targets.size() == 3, "compute runtime targets"); + expect(compute_runtime.targets[0].target.name == "cuda:0", "compute target0"); + expect(compute_runtime.targets[1].target.name == "hip:0", "compute target1"); + expect(compute_runtime.targets[2].target.backend == "cpu", "compute target2"); + expect(compute_runtime.targets[0].placement.total_hot == 4, "compute target0 hot"); + expect(compute_runtime.targets[1].placement.total_hot == 2, "compute target1 hot"); + expect(compute_runtime.targets[2].placement.total_hot == 2, "compute target2 hot"); + expect(compute_runtime.targets[0].placement.hot_expert_ids[0][0] == 0, "compute t0 l0 e0"); + expect(compute_runtime.targets[1].placement.hot_expert_ids[1][0] == 0, "compute t1 l1 e0"); + expect(compute_runtime.targets[1].placement.hot_expert_ids[1][1] == 1, "compute t1 l1 e1"); + expect(compute_runtime.targets[2].placement.hot_expert_ids[1][0] == 2, "compute t2 l1 e2"); + expect(compute_runtime.targets[2].placement.hot_expert_ids[1][1] == 3, "compute t2 l1 e3"); + expect(compute_runtime.target_index(1, 2) == 2, "compute target lookup l1 e2"); + expect(compute_runtime.local_index(1, 2) == 0, "compute local lookup l1 e2"); + expect(compute_runtime.target_index(1, 3) == 2, "compute target lookup l1 e3"); + expect(compute_runtime.local_index(1, 3) == 1, "compute local lookup l1 e3"); + expect(compute_runtime.target_index(-1, 0) == -1, "compute target lookup rejects layer"); + expect(compute_runtime.local_index(0, 4) == -1, "compute local lookup rejects expert"); +} + +static void test_compute_runtime_preserves_target_local_order() { + const ExpertSplitPlan plan = build_sample_plan(); + ExpertSplitRuntime runtime; + std::string err; + expect(build_expert_split_runtime(plan, runtime, &err), err.c_str()); + + ExpertSplitComputeRuntime compute_runtime; + expect(build_expert_split_compute_runtime(runtime, /*n_expert_used=*/2, + compute_runtime, &err), + err.c_str()); + + const auto & hip_layer = compute_runtime.targets[1].placement.hot_expert_ids[1]; + const auto & cpu_layer = compute_runtime.targets[2].placement.hot_expert_ids[1]; + expect(hip_layer.size() == 2, "hip layer count"); + expect(cpu_layer.size() == 2, "cpu layer count"); + expect(hip_layer[0] == 0 && hip_layer[1] == 1, "hip local order follows score"); + expect(cpu_layer[0] == 2 && cpu_layer[1] == 3, "cpu local order follows score"); +} + +static ExpertSplitPlan build_reordered_target_plan() { + ExpertSplitPlan plan; + plan.n_layer = 1; + plan.n_expert = 4; + plan.targets = { + {"cuda:0", "cuda", 0, 2, 0, false}, + {"hip:0", "hip", 0, 2, 0, false}, + }; + plan.target_used_bytes = {2, 2}; + plan.assignments = { + {0, 1, 10.0}, + {1, 1, 9.0}, + {0, 1, 8.0}, + {1, 1, 7.0}, + }; + return plan; +} + +static void test_compute_runtime_fingerprint_changes_with_assignment() { + ExpertSplitRuntime runtime_a; + ExpertSplitRuntime runtime_b; + std::string err; + + expect(build_expert_split_runtime(build_sample_plan(), runtime_a, &err), err.c_str()); + expect(build_expert_split_runtime(build_reordered_target_plan(), runtime_b, &err), err.c_str()); + + ExpertSplitComputeRuntime compute_a; + ExpertSplitComputeRuntime compute_b; + expect(build_expert_split_compute_runtime(runtime_a, /*n_expert_used=*/2, + compute_a, &err), + err.c_str()); + expect(build_expert_split_compute_runtime(runtime_b, /*n_expert_used=*/2, + compute_b, &err), + err.c_str()); + + const uint64_t fp_a = expert_split_compute_runtime_fingerprint(compute_a); + const uint64_t fp_b = expert_split_compute_runtime_fingerprint(compute_b); + expect(fp_a != fp_b, "compute runtime fingerprint changes when placement changes"); +} + +static void test_runtime_rejects_invalid_plan() { + ExpertSplitPlan plan; + plan.n_layer = 1; + plan.n_expert = 1; + plan.targets = {{"cpu", "cpu", -1, 0, 0, true}}; + plan.target_used_bytes = {0}; + plan.assignments = {{-1, 0, 0.0}}; + + ExpertSplitRuntime runtime; + std::string err; + expect(!build_expert_split_runtime(plan, runtime, &err), "invalid plan rejected"); + expect(!err.empty(), "invalid plan error"); +} + +int main() { + test_runtime_materializes_targets_and_cpu_fallback(); + test_runtime_local_order_is_score_stable(); + test_target_placement_roundtrip(); + test_all_target_placements_materialize(); + test_compute_runtime_materializes_target_views(); + test_compute_runtime_preserves_target_local_order(); + test_compute_runtime_fingerprint_changes_with_assignment(); + test_runtime_rejects_invalid_plan(); + std::printf("OK\n"); + return 0; +} diff --git a/server/test/test_expert_split_state.cpp b/server/test/test_expert_split_state.cpp new file mode 100644 index 000000000..77f2f88cb --- /dev/null +++ b/server/test/test_expert_split_state.cpp @@ -0,0 +1,86 @@ +#include "../src/common/expert_split_state.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +static void expect(bool cond, const char * msg) { + if (!cond) { + std::fprintf(stderr, "FAIL: %s\n", msg); + std::exit(1); + } +} + +static void test_builds_shared_expert_split_state_components() { + ExpertSplitConfig cfg; + cfg.n_layer = 2; + cfg.n_expert = 3; + + const std::vector targets = { + {"cuda:0", "cuda", 0, 3, 0, false}, + {"hip:0", "hip", 0, 1, 0, false}, + }; + const std::vector units = { + {0, 0, 1, 100.0}, + {0, 1, 1, 90.0}, + {0, 2, 1, 80.0}, + {1, 0, 1, 70.0}, + {1, 1, 1, 60.0}, + {1, 2, 1, 50.0}, + }; + + ExpertSplitStateComponents state; + std::string err; + expect(build_expert_split_state(cfg, targets, units, /*n_expert_used=*/2, + state, &err), + err.c_str()); + expect(state.matches(2, 3, 2), "shared state matches dims"); + expect(state.plan.at(0, 0).target_index == 0, "plan target lookup"); + expect(state.runtime.target_index(1, 2) == 2, "runtime includes cpu fallback"); + expect(state.compute_runtime.local_index(1, 2) == 1, "compute runtime local id"); + expect(state.materialization.ordered_cold_union, "materialization ordered cold union"); +} + +static void test_builds_generic_layer_mapping() { + ExpertSplitLayerMapping mapping; + std::string err; + expect(build_expert_split_layer_mapping( + /*n_total_layer=*/5, {1, 3}, mapping, &err), + err.c_str()); + expect(mapping.n_total_layer == 5, "mapping total layers"); + expect(mapping.split_layer_count() == 2, "mapping split count"); + expect(mapping.physical_layer_by_split_layer[0] == 1, "mapping split0"); + expect(mapping.physical_layer_by_split_layer[1] == 3, "mapping split1"); + expect(mapping.split_layer_for_physical(1) == 0, "mapping reverse 1"); + expect(mapping.split_layer_for_physical(3) == 1, "mapping reverse 3"); + expect(mapping.split_layer_for_physical(0) == -1, "mapping dense lead"); + expect(mapping.is_split_layer(3), "mapping is split layer"); + expect(!mapping.is_split_layer(4), "mapping rejects dense layer"); +} + +static void test_builds_contiguous_layer_mapping() { + ExpertSplitLayerMapping mapping; + std::string err; + expect(build_contiguous_expert_split_layer_mapping( + /*n_total_layer=*/4, /*first_split_layer=*/1, + /*n_split_layer=*/3, mapping, &err), + err.c_str()); + expect(mapping.physical_layer_by_split_layer.size() == 3, "contiguous mapping size"); + expect(mapping.physical_layer_by_split_layer[0] == 1, "contiguous split0"); + expect(mapping.physical_layer_by_split_layer[1] == 2, "contiguous split1"); + expect(mapping.physical_layer_by_split_layer[2] == 3, "contiguous split2"); + expect(mapping.split_layer_for_physical(1) == 0, "contiguous reverse1"); + expect(mapping.split_layer_for_physical(2) == 1, "contiguous reverse2"); + expect(mapping.split_layer_for_physical(3) == 2, "contiguous reverse3"); +} + +int main() { + test_builds_shared_expert_split_state_components(); + test_builds_generic_layer_mapping(); + test_builds_contiguous_layer_mapping(); + std::printf("OK\n"); + return 0; +} diff --git a/server/test/test_expert_split_target_config.cpp b/server/test/test_expert_split_target_config.cpp new file mode 100644 index 000000000..523c2b24a --- /dev/null +++ b/server/test/test_expert_split_target_config.cpp @@ -0,0 +1,113 @@ +#include "../src/common/expert_split_target_config.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +static void expect(bool cond, const char * msg) { + if (!cond) { + std::fprintf(stderr, "FAIL: %s\n", msg); + std::exit(1); + } +} + +static void test_parse_target_list() { + std::vector specs; + std::string err; + expect(parse_expert_split_target_list("cuda:0,cuda:1,cpu", specs, &err), err.c_str()); + expect(specs.size() == 3, "target count"); + expect(specs[0].backend == PlacementBackend::Cuda, "target0 backend"); + expect(specs[0].device_id == 0, "target0 device"); + expect(specs[1].backend == PlacementBackend::Cuda, "target1 backend"); + expect(specs[1].device_id == 1, "target1 device"); + expect(specs[2].unlimited, "cpu target unlimited"); +} + +static void test_parse_capacity_overrides() { + std::vector caps; + std::string err; + expect(parse_expert_split_capacity_overrides("1024MB,auto,2GB", caps, &err), err.c_str()); + expect(caps.size() == 3, "cap count"); + expect(caps[0] == 1024ULL * 1024ULL * 1024ULL, "cap0 bytes"); + expect(caps[1] == 0, "cap1 auto"); + expect(caps[2] == 2ULL * 1024ULL * 1024ULL * 1024ULL, "cap2 bytes"); +} + +static void test_build_targets_with_primary_override() { + std::vector specs; + std::string err; + expect(parse_expert_split_target_list("cuda:0,cuda:1,cpu", specs, &err), err.c_str()); + specs[1].auto_capacity = false; + specs[1].capacity_bytes = 3ULL * 1024ULL * 1024ULL * 1024ULL; + + std::vector targets; + expect(build_expert_split_targets(specs, + /*primary_capacity_bytes=*/5ULL * 1024ULL * 1024ULL * 1024ULL, + targets, &err), + err.c_str()); + expect(targets.size() == 3, "built target count"); + expect(targets[0].name == "cuda:0", "built target0 name"); + expect(targets[0].capacity_bytes == 5ULL * 1024ULL * 1024ULL * 1024ULL, "primary cap"); + expect(targets[1].name == "cuda:1", "built target1 name"); + expect(targets[1].capacity_bytes == 3ULL * 1024ULL * 1024ULL * 1024ULL, "secondary cap"); + expect(targets[2].backend == "cpu", "cpu backend"); + expect(targets[2].unlimited, "cpu unlimited"); +} + +static void test_reject_duplicate_targets() { + std::vector specs; + std::string err; + expect(!parse_expert_split_target_list("cuda:0,cuda:0", specs, &err), "duplicate rejected"); + expect(!err.empty(), "duplicate error"); +} + +static void test_resolve_targets_from_env() { + std::string err; + ::setenv("DFLASH_TEST_EXPERT_TARGETS", "cuda:0,cpu", 1); + ::setenv("DFLASH_TEST_EXPERT_TARGET_CAPS", "4GB,auto", 1); + + std::vector targets; + expect(resolve_expert_split_targets_from_env( + "DFLASH_TEST_EXPERT_TARGETS", + "DFLASH_TEST_EXPERT_TARGET_CAPS", + /*primary_capacity_bytes=*/6ULL * 1024ULL * 1024ULL * 1024ULL, + targets, &err), + err.c_str()); + expect(targets.size() == 2, "resolved target count"); + expect(targets[0].name == "cuda:0", "resolved target0 name"); + expect(targets[0].capacity_bytes == 6ULL * 1024ULL * 1024ULL * 1024ULL, + "resolved primary cap"); + expect(targets[1].backend == "cpu", "resolved cpu backend"); + expect(targets[1].unlimited, "resolved cpu unlimited"); + + ::unsetenv("DFLASH_TEST_EXPERT_TARGETS"); + ::unsetenv("DFLASH_TEST_EXPERT_TARGET_CAPS"); +} + +static void test_validate_primary_target() { + std::vector targets = { + {"cuda:0", "cuda", 0, 1, 0, false}, + {"cpu", "cpu", -1, 0, 0, true}, + }; + std::string err; + expect(validate_primary_expert_split_target( + targets, PlacementBackend::Cuda, 0, &err), + err.c_str()); + expect(!validate_primary_expert_split_target( + targets, PlacementBackend::Cuda, 1, &err), + "primary target mismatch rejected"); +} + +int main() { + test_parse_target_list(); + test_parse_capacity_overrides(); + test_build_targets_with_primary_override(); + test_reject_duplicate_targets(); + test_resolve_targets_from_env(); + test_validate_primary_target(); + std::printf("OK\n"); + return 0; +} diff --git a/server/test/test_moe_expert_compute_multi_target.cpp b/server/test/test_moe_expert_compute_multi_target.cpp new file mode 100644 index 000000000..81f99082c --- /dev/null +++ b/server/test/test_moe_expert_compute_multi_target.cpp @@ -0,0 +1,354 @@ +#include "../src/common/moe_expert_compute.h" + +#include +#include +#include +#include +#include + +using namespace dflash::common; + +static void expect(bool cond, const char * msg) { + if (!cond) { + std::fprintf(stderr, "FAIL: %s\n", msg); + std::exit(1); + } +} + +static void expect_close(float got, float want, const char * msg) { + if (std::fabs(got - want) > 1e-5f) { + std::fprintf(stderr, "FAIL: %s got=%f want=%f\n", msg, got, want); + std::exit(1); + } +} + +struct FakeMoeExpertCompute final : MoeExpertCompute { + explicit FakeMoeExpertCompute(int target_id_) : target_id(target_id_) {} + + bool prefers_padded_batch() const override { + return padded_preferred; + } + + bool compute(const MoeExpertLayer &, + const float *, + const int32_t *, + const float *, + int, + int, + int, + float *) override { + ++single_calls; + return false; + } + + bool compute_batch(const MoeExpertLayer &, + const float * input, + const int32_t * ids, + const float * weights, + int n_tokens, + int n_selected, + int n_embd, + int, + float * output) override { + ++batch_calls; + batch_tokens.push_back(n_tokens); + batch_selected.push_back(n_selected); + for (int t = 0; t < n_tokens; ++t) { + float acc = 0.0f; + for (int i = 0; i < n_selected; ++i) { + const size_t idx = (size_t)t * (size_t)n_selected + (size_t)i; + acc += weights[idx] * (float)(target_id * 100 + ids[idx]); + } + for (int e = 0; e < n_embd; ++e) { + output[(size_t)t * (size_t)n_embd + (size_t)e] = + input[(size_t)t * (size_t)n_embd + (size_t)e] + acc + (float)e; + } + } + return true; + } + + int target_id = 0; + int single_calls = 0; + int batch_calls = 0; + bool padded_preferred = false; + std::vector batch_tokens; + std::vector batch_selected; +}; + +static void test_multi_target_compute_batch_groups_by_target_and_count() { + MoeMultiTargetExpertRuntime runtime; + runtime.targets.resize(2); + runtime.layer_routes.resize(1); + runtime.layers.resize(1); + runtime.layers[0].layer_idx = 0; + runtime.layers[0].cold_global_by_local = {0, 1, 2, 3}; + runtime.enabled = true; + + auto target0_compute = std::make_unique(1); + auto target1_compute = std::make_unique(2); + FakeMoeExpertCompute * target0 = target0_compute.get(); + FakeMoeExpertCompute * target1 = target1_compute.get(); + + runtime.targets[0].target_index = 0; + runtime.targets[0].placement.total_hot = 2; + runtime.targets[0].compute_active = true; + runtime.targets[0].runtime.compute = std::move(target0_compute); + runtime.targets[0].runtime.layers.resize(1); + runtime.targets[0].runtime.layers[0].layer_idx = 0; + runtime.targets[0].runtime.layers[0].cold_global_by_local = {0, 1}; + + runtime.targets[1].target_index = 1; + runtime.targets[1].placement.total_hot = 2; + runtime.targets[1].compute_active = true; + runtime.targets[1].runtime.compute = std::move(target1_compute); + runtime.targets[1].runtime.layers.resize(1); + runtime.targets[1].runtime.layers[0].layer_idx = 0; + runtime.targets[1].runtime.layers[0].cold_global_by_local = {2, 3}; + + auto & routes = runtime.layer_routes[0].route_by_union_local; + routes.resize(4); + routes[0] = {0, 0, 0}; + routes[1] = {0, 1, 1}; + routes[2] = {1, 0, 2}; + routes[3] = {1, 1, 3}; + + runtime.compute = make_multi_target_moe_expert_compute(&runtime); + + const int n_tokens = 3; + const int n_selected = 3; + const int n_embd = 2; + const float input[] = { + 1.0f, 2.0f, + 3.0f, 4.0f, + 5.0f, 6.0f, + }; + const int32_t ids[] = { + 0, 2, 3, + 1, 0, 2, + 2, 3, 1, + }; + const float weights[] = { + 1.0f, 2.0f, 3.0f, + 4.0f, 5.0f, 6.0f, + 7.0f, 8.0f, 9.0f, + }; + float output[n_tokens * n_embd] = {}; + + expect(runtime.compute->compute_batch(runtime.layers[0], input, ids, weights, + n_tokens, n_selected, n_embd, + /*n_ff=*/4, output), + "multi-target compute_batch"); + + expect(target0->single_calls == 0, "target0 did not use single compute"); + expect(target1->single_calls == 0, "target1 did not use single compute"); + expect(target0->batch_calls == 2, "target0 grouped into two batch calls"); + expect(target1->batch_calls == 2, "target1 grouped into two batch calls"); + + expect(target0->batch_selected[0] == 1, "target0 first selected count"); + expect(target0->batch_selected[1] == 2, "target0 second selected count"); + expect(target1->batch_selected[0] == 1, "target1 first selected count"); + expect(target1->batch_selected[1] == 2, "target1 second selected count"); + expect(target0->batch_tokens[0] + target0->batch_tokens[1] == n_tokens, + "target0 covers all tokens"); + expect(target1->batch_tokens[0] + target1->batch_tokens[1] == n_tokens, + "target1 covers all tokens"); + + const float expected[] = { + 1105.0f, 1109.0f, + 2110.0f, 2114.0f, + 3927.0f, 3931.0f, + }; + for (int i = 0; i < n_tokens * n_embd; ++i) { + expect_close(output[i], expected[i], "accumulated output"); + } +} + +static void test_multi_target_compute_batch_accumulates_after_direct_write() { + MoeMultiTargetExpertRuntime runtime; + runtime.targets.resize(2); + runtime.layer_routes.resize(1); + runtime.layers.resize(1); + runtime.layers[0].layer_idx = 0; + runtime.layers[0].cold_global_by_local = {0, 1}; + runtime.enabled = true; + + auto target0_compute = std::make_unique(1); + auto target1_compute = std::make_unique(2); + + runtime.targets[0].target_index = 0; + runtime.targets[0].placement.total_hot = 1; + runtime.targets[0].compute_active = true; + runtime.targets[0].runtime.compute = std::move(target0_compute); + runtime.targets[0].runtime.layers.resize(1); + runtime.targets[0].runtime.layers[0].layer_idx = 0; + runtime.targets[0].runtime.layers[0].cold_global_by_local = {0}; + + runtime.targets[1].target_index = 1; + runtime.targets[1].placement.total_hot = 1; + runtime.targets[1].compute_active = true; + runtime.targets[1].runtime.compute = std::move(target1_compute); + runtime.targets[1].runtime.layers.resize(1); + runtime.targets[1].runtime.layers[0].layer_idx = 0; + runtime.targets[1].runtime.layers[0].cold_global_by_local = {1}; + + auto & routes = runtime.layer_routes[0].route_by_union_local; + routes.resize(2); + routes[0] = {0, 0, 0}; + routes[1] = {1, 0, 1}; + + runtime.compute = make_multi_target_moe_expert_compute(&runtime); + + const int n_tokens = 2; + const int n_selected = 2; + const int n_embd = 2; + const float input[] = { + 1.0f, 2.0f, + 3.0f, 4.0f, + }; + const int32_t ids[] = { + 0, 1, + 0, 1, + }; + const float weights[] = { + 1.0f, 2.0f, + 3.0f, 4.0f, + }; + float output[n_tokens * n_embd] = {}; + + expect(runtime.compute->compute_batch(runtime.layers[0], input, ids, weights, + n_tokens, n_selected, n_embd, + /*n_ff=*/4, output), + "multi-target direct-write accumulation"); + + const float expected[] = { + 502.0f, 506.0f, + 1106.0f, 1110.0f, + }; + for (int i = 0; i < n_tokens * n_embd; ++i) { + expect_close(output[i], expected[i], "direct write then add output"); + } +} + +static void test_multi_target_compute_batch_reuses_scratch_without_stale_groups() { + MoeMultiTargetExpertRuntime runtime; + runtime.targets.resize(2); + runtime.layer_routes.resize(1); + runtime.layers.resize(1); + runtime.layers[0].layer_idx = 0; + runtime.layers[0].cold_global_by_local = {0, 1, 2}; + runtime.enabled = true; + + auto target0_compute = std::make_unique(1); + auto target1_compute = std::make_unique(2); + FakeMoeExpertCompute * target0 = target0_compute.get(); + FakeMoeExpertCompute * target1 = target1_compute.get(); + + runtime.targets[0].target_index = 0; + runtime.targets[0].placement.total_hot = 2; + runtime.targets[0].compute_active = true; + runtime.targets[0].runtime.compute = std::move(target0_compute); + runtime.targets[0].runtime.layers.resize(1); + runtime.targets[0].runtime.layers[0].layer_idx = 0; + runtime.targets[0].runtime.layers[0].cold_global_by_local = {0, 1}; + + runtime.targets[1].target_index = 1; + runtime.targets[1].placement.total_hot = 1; + runtime.targets[1].compute_active = true; + runtime.targets[1].runtime.compute = std::move(target1_compute); + runtime.targets[1].runtime.layers.resize(1); + runtime.targets[1].runtime.layers[0].layer_idx = 0; + runtime.targets[1].runtime.layers[0].cold_global_by_local = {2}; + + auto & routes = runtime.layer_routes[0].route_by_union_local; + routes.resize(3); + routes[0] = {0, 0, 0}; + routes[1] = {0, 1, 1}; + routes[2] = {1, 0, 2}; + + runtime.compute = make_multi_target_moe_expert_compute(&runtime); + + const int n_embd = 1; + const float input1[] = {1.0f, 2.0f}; + const int32_t ids1[] = { + 0, 1, 2, + 0, 1, 2, + }; + const float weights1[] = { + 1.0f, 1.0f, 1.0f, + 1.0f, 1.0f, 1.0f, + }; + float output1[2] = {}; + expect(runtime.compute->compute_batch(runtime.layers[0], input1, ids1, weights1, + /*n_tokens=*/2, /*n_selected=*/3, n_embd, + /*n_ff=*/4, output1), + "multi-target first scratch reuse batch"); + + const float input2[] = {3.0f}; + const int32_t ids2[] = {2, 0}; + const float weights2[] = {4.0f, 5.0f}; + float output2[1] = {}; + expect(runtime.compute->compute_batch(runtime.layers[0], input2, ids2, weights2, + /*n_tokens=*/1, /*n_selected=*/2, n_embd, + /*n_ff=*/4, output2), + "multi-target second scratch reuse batch"); + + expect(target0->batch_calls == 2, "target0 called once per batch"); + expect(target1->batch_calls == 2, "target1 called once per batch"); + expect(target0->batch_selected.back() == 1, "target0 second batch selected"); + expect(target1->batch_selected.back() == 1, "target1 second batch selected"); + expect(target0->batch_tokens.back() == 1, "target0 second batch tokens"); + expect(target1->batch_tokens.back() == 1, "target1 second batch tokens"); + expect_close(output2[0], 1306.0f, "second batch output excludes stale buckets"); +} + +static void test_multi_target_prefers_grouped_batches_with_multiple_active_targets() { + MoeMultiTargetExpertRuntime runtime; + runtime.targets.resize(2); + runtime.enabled = true; + + auto target0_compute = std::make_unique(1); + auto target1_compute = std::make_unique(2); + target1_compute->padded_preferred = true; + + runtime.targets[0].compute_active = true; + runtime.targets[0].runtime.compute = std::move(target0_compute); + runtime.targets[1].compute_active = true; + runtime.targets[1].runtime.compute = std::move(target1_compute); + runtime.compute = make_multi_target_moe_expert_compute(&runtime); + + expect(!runtime.compute->prefers_padded_batch(), + "multi-target disables padded batch preference across multiple active targets"); +} + +static void test_multi_target_inactive_primary_does_not_require_compute() { + MoeMultiTargetExpertRuntime runtime; + runtime.targets.resize(2); + runtime.enabled = true; + + runtime.targets[0].target_index = 0; + runtime.targets[0].placement.total_hot = 4; + runtime.targets[0].compute_active = false; + + auto target1_compute = std::make_unique(2); + target1_compute->padded_preferred = true; + runtime.targets[1].target_index = 1; + runtime.targets[1].placement.total_hot = 2; + runtime.targets[1].compute_active = true; + runtime.targets[1].runtime.compute = std::move(target1_compute); + + runtime.compute = make_multi_target_moe_expert_compute(&runtime); + expect(runtime.compute->healthy(), + "inactive primary target does not make runtime unhealthy"); + expect(runtime.compute->prefers_padded_batch(), + "inactive primary target is skipped for preferences"); +} + +int main() { + test_multi_target_compute_batch_groups_by_target_and_count(); + test_multi_target_compute_batch_accumulates_after_direct_write(); + test_multi_target_compute_batch_reuses_scratch_without_stale_groups(); + test_multi_target_prefers_grouped_batches_with_multiple_active_targets(); + test_multi_target_inactive_primary_does_not_require_compute(); + std::printf("OK\n"); + return 0; +} diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 617242eb6..a39785ada 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -20,6 +20,7 @@ #include "common/sampler.h" #include "common/backend_precision.h" #include "common/backend_ipc.h" +#include "common/moe_expert_compute.h" #include "placement/pflash_placement.h" #include "common/io_utils.h" #include "placement/placement_config.h" @@ -3040,6 +3041,44 @@ static void test_backend_ipc_shared_payload_map_sizing() { std::numeric_limits::max(), map_bytes)); } +static void test_moe_expert_compute_prepare_batch_default_is_opt_in() { + unsetenv("DFLASH_MOE_EXPERT_COMPUTE_PREPARE_BATCH"); + unsetenv("DFLASH_MOE_EXPERT_COMPUTE_PREPARE_SELECTED"); + TEST_ASSERT(moe_expert_compute_prepare_batch_limit_from_env() == 0); + TEST_ASSERT(moe_expert_compute_prepare_selected_limit_from_env() == 0); + + setenv("DFLASH_MOE_EXPERT_COMPUTE_PREPARE_BATCH", "8", 1); + setenv("DFLASH_MOE_EXPERT_COMPUTE_PREPARE_SELECTED", "3", 1); + TEST_ASSERT(moe_expert_compute_prepare_batch_limit_from_env() == 8); + TEST_ASSERT(moe_expert_compute_prepare_selected_limit_from_env() == 3); + + setenv("DFLASH_MOE_EXPERT_COMPUTE_PREPARE_BATCH", "99999", 1); + setenv("DFLASH_MOE_EXPERT_COMPUTE_PREPARE_SELECTED", "99999", 1); + TEST_ASSERT(moe_expert_compute_prepare_batch_limit_from_env() == 4096); + TEST_ASSERT(moe_expert_compute_prepare_selected_limit_from_env() == 4096); + + setenv("DFLASH_MOE_EXPERT_COMPUTE_PREPARE_BATCH", "0", 1); + setenv("DFLASH_MOE_EXPERT_COMPUTE_PREPARE_SELECTED", "0", 1); + TEST_ASSERT(moe_expert_compute_prepare_batch_limit_from_env() == 0); + TEST_ASSERT(moe_expert_compute_prepare_selected_limit_from_env() == 0); +} + +static void test_moe_hybrid_all_routed_are_hot_handles_expert_ids_over_255() { + MoeHybridLayerStorage storage; + storage.hot_local_by_global.assign(512, -1); + + storage.hot_local_by_global[42] = 0; + storage.hot_local_by_global[300] = 1; + + const int32_t all_hot[] = {42, 300}; + TEST_ASSERT_MSG(storage.all_routed_are_hot(all_hot, 2), + "residency check must treat >255 expert ids as hot when mapped"); + + storage.hot_local_by_global[300] = -1; + TEST_ASSERT_MSG(!storage.all_routed_are_hot(all_hot, 2), + "residency check must reject >255 expert ids when not mapped hot"); +} + // ═══════════════════════════════════════════════════════════════════════ // Sampler tests (model-independent, CPU-only) // ═══════════════════════════════════════════════════════════════════════ @@ -4419,6 +4458,8 @@ int main() { RUN_TEST(test_backend_ipc_payload_transport_parse); RUN_TEST(test_backend_ipc_payload_bounds); RUN_TEST(test_backend_ipc_shared_payload_map_sizing); + RUN_TEST(test_moe_expert_compute_prepare_batch_default_is_opt_in); + RUN_TEST(test_moe_hybrid_all_routed_are_hot_handles_expert_ids_over_255); std::fprintf(stderr, "\n── Jinja chat template ──\n"); RUN_TEST(test_jinja_render_basic); From abd400a999a8944f72b35b0b736b7a496b256be6 Mon Sep 17 00:00:00 2001 From: weicj Date: Fri, 3 Jul 2026 01:40:27 +0800 Subject: [PATCH 2/2] test(server): trim expert-split test surface --- server/CMakeLists.txt | 16 --- .../test_expert_split_materialization.cpp | 104 -------------- server/test/test_expert_split_runtime.cpp | 129 ++++++++++++++++++ server/test/test_expert_split_state.cpp | 86 ------------ 4 files changed, 129 insertions(+), 206 deletions(-) delete mode 100644 server/test/test_expert_split_materialization.cpp delete mode 100644 server/test/test_expert_split_state.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 4caffc41f..6101ec663 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -753,14 +753,6 @@ if(DFLASH27B_TESTS) endif() target_link_libraries(test_expert_split_runtime PRIVATE dflash_common) endif() - if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_expert_split_state.cpp") - add_executable(test_expert_split_state test/test_expert_split_state.cpp) - target_include_directories(test_expert_split_state PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) - if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") - target_include_directories(test_expert_split_state PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) - endif() - target_link_libraries(test_expert_split_state PRIVATE dflash_common) - endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_expert_split_target_config.cpp") add_executable(test_expert_split_target_config test/test_expert_split_target_config.cpp) target_include_directories(test_expert_split_target_config PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) @@ -769,14 +761,6 @@ if(DFLASH27B_TESTS) endif() target_link_libraries(test_expert_split_target_config PRIVATE dflash_common) endif() - if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_expert_split_materialization.cpp") - add_executable(test_expert_split_materialization test/test_expert_split_materialization.cpp) - target_include_directories(test_expert_split_materialization PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) - if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") - target_include_directories(test_expert_split_materialization PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) - endif() - target_link_libraries(test_expert_split_materialization PRIVATE dflash_common) - endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_moe_expert_compute_multi_target.cpp") add_executable(test_moe_expert_compute_multi_target test/test_moe_expert_compute_multi_target.cpp) target_include_directories(test_moe_expert_compute_multi_target PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) diff --git a/server/test/test_expert_split_materialization.cpp b/server/test/test_expert_split_materialization.cpp deleted file mode 100644 index 7823aaa98..000000000 --- a/server/test/test_expert_split_materialization.cpp +++ /dev/null @@ -1,104 +0,0 @@ -#include "../src/common/expert_split_materialization.h" -#include "../src/common/expert_split_plan.h" -#include "../src/common/expert_split_runtime.h" - -#include -#include -#include -#include - -using namespace dflash::common; - -static void expect(bool cond, const char * msg) { - if (!cond) { - std::fprintf(stderr, "FAIL: %s\n", msg); - std::exit(1); - } -} - -static ExpertSplitPlan build_multi_target_plan() { - ExpertSplitConfig cfg; - cfg.n_layer = 2; - cfg.n_expert = 4; - - const std::vector targets = { - {"cuda:0", "cuda", 0, 4, 0, false}, - {"hip:0", "hip", 0, 2, 0, false}, - }; - - const std::vector units = { - {0, 0, 1, 100.0}, {0, 1, 1, 90.0}, {0, 2, 1, 80.0}, {0, 3, 1, 70.0}, - {1, 0, 1, 60.0}, {1, 1, 1, 50.0}, {1, 2, 1, 40.0}, {1, 3, 1, 30.0}, - }; - - ExpertSplitPlan plan; - std::string err; - expect(build_expert_split_plan(cfg, targets, units, plan, &err), err.c_str()); - return plan; -} - -static void test_materialization_builds_ordered_cold_union() { - const ExpertSplitPlan plan = build_multi_target_plan(); - ExpertSplitRuntime runtime; - std::string err; - expect(build_expert_split_runtime(plan, runtime, &err), err.c_str()); - - ExpertSplitMaterialization materialization; - expect(build_expert_split_materialization(runtime, /*n_expert_used=*/2, - materialization, &err), - err.c_str()); - - expect(materialization.matches(2, 4, 2), "materialization dims"); - expect(materialization.ordered_cold_union, "ordered union enabled"); - expect(materialization.targets.size() == 3, "targets include cpu fallback"); - expect(materialization.primary_placement.total_hot == 4, "primary hot total"); - - const auto & l0_hot = materialization.primary_placement.hot_expert_ids[0]; - const auto & l1_hot = materialization.primary_placement.hot_expert_ids[1]; - expect(l0_hot.size() == 4, "layer0 primary hot count"); - expect(l1_hot.empty(), "layer1 primary hot empty"); - expect(l0_hot[0] == 0 && l0_hot[1] == 1 && l0_hot[2] == 2 && l0_hot[3] == 3, - "layer0 primary hot order"); - - const auto & l0_cold = materialization.cold_expert_ids_by_layer[0]; - const auto & l1_cold = materialization.cold_expert_ids_by_layer[1]; - expect(l0_cold.empty(), "layer0 cold empty"); - expect(l1_cold.size() == 4, "layer1 cold count"); - expect(l1_cold[0] == 0 && l1_cold[1] == 1 && l1_cold[2] == 2 && - l1_cold[3] == 3, "layer1 cold preserves target order"); -} - -static void test_materialization_single_explicit_gpu_keeps_single_primary_shape() { - ExpertSplitConfig cfg; - cfg.n_layer = 1; - cfg.n_expert = 4; - - const std::vector targets = { - {"cuda:0", "cuda", 0, 2, 0, false}, - }; - const std::vector units = { - {0, 0, 1, 40.0}, {0, 1, 1, 30.0}, {0, 2, 1, 20.0}, {0, 3, 1, 10.0}, - }; - - ExpertSplitPlan plan; - std::string err; - expect(build_expert_split_plan(cfg, targets, units, plan, &err), err.c_str()); - - ExpertSplitRuntime runtime; - expect(build_expert_split_runtime(plan, runtime, &err), err.c_str()); - - ExpertSplitMaterialization materialization; - expect(build_expert_split_materialization(runtime, /*n_expert_used=*/2, - materialization, &err), - err.c_str()); - expect(materialization.matches(1, 4, 2), "single-primary materialization dims"); - expect(!materialization.ordered_cold_union, "single-primary path keeps existing cold order"); - expect(materialization.cold_expert_ids_by_layer.empty(), "no cold override"); -} - -int main() { - test_materialization_builds_ordered_cold_union(); - test_materialization_single_explicit_gpu_keeps_single_primary_shape(); - std::printf("OK\n"); - return 0; -} diff --git a/server/test/test_expert_split_runtime.cpp b/server/test/test_expert_split_runtime.cpp index cfa7c6213..943473b38 100644 --- a/server/test/test_expert_split_runtime.cpp +++ b/server/test/test_expert_split_runtime.cpp @@ -1,6 +1,8 @@ #include "../src/common/expert_split_plan.h" #include "../src/common/expert_split_compute_runtime.h" +#include "../src/common/expert_split_materialization.h" #include "../src/common/expert_split_runtime.h" +#include "../src/common/expert_split_state.h" #include #include @@ -229,6 +231,128 @@ static void test_compute_runtime_fingerprint_changes_with_assignment() { expect(fp_a != fp_b, "compute runtime fingerprint changes when placement changes"); } +static void test_materialization_builds_ordered_cold_union() { + const ExpertSplitPlan plan = build_sample_plan(); + ExpertSplitRuntime runtime; + std::string err; + expect(build_expert_split_runtime(plan, runtime, &err), err.c_str()); + + ExpertSplitMaterialization materialization; + expect(build_expert_split_materialization(runtime, /*n_expert_used=*/2, + materialization, &err), + err.c_str()); + + expect(materialization.matches(2, 4, 2), "materialization dims"); + expect(materialization.ordered_cold_union, "ordered union enabled"); + expect(materialization.targets.size() == 3, "targets include cpu fallback"); + expect(materialization.primary_placement.total_hot == 4, "primary hot total"); + + const auto & l0_hot = materialization.primary_placement.hot_expert_ids[0]; + const auto & l1_hot = materialization.primary_placement.hot_expert_ids[1]; + expect(l0_hot.size() == 4, "layer0 primary hot count"); + expect(l1_hot.empty(), "layer1 primary hot empty"); + expect(l0_hot[0] == 0 && l0_hot[1] == 1 && l0_hot[2] == 2 && l0_hot[3] == 3, + "layer0 primary hot order"); + + const auto & l0_cold = materialization.cold_expert_ids_by_layer[0]; + const auto & l1_cold = materialization.cold_expert_ids_by_layer[1]; + expect(l0_cold.empty(), "layer0 cold empty"); + expect(l1_cold.size() == 4, "layer1 cold count"); + expect(l1_cold[0] == 0 && l1_cold[1] == 1 && l1_cold[2] == 2 && + l1_cold[3] == 3, "layer1 cold preserves target order"); +} + +static void test_materialization_single_explicit_gpu_keeps_single_primary_shape() { + ExpertSplitConfig cfg; + cfg.n_layer = 1; + cfg.n_expert = 4; + + const std::vector targets = { + {"cuda:0", "cuda", 0, 2, 0, false}, + }; + const std::vector units = { + {0, 0, 1, 40.0}, {0, 1, 1, 30.0}, {0, 2, 1, 20.0}, {0, 3, 1, 10.0}, + }; + + ExpertSplitPlan plan; + std::string err; + expect(build_expert_split_plan(cfg, targets, units, plan, &err), err.c_str()); + + ExpertSplitRuntime runtime; + expect(build_expert_split_runtime(plan, runtime, &err), err.c_str()); + + ExpertSplitMaterialization materialization; + expect(build_expert_split_materialization(runtime, /*n_expert_used=*/2, + materialization, &err), + err.c_str()); + expect(materialization.matches(1, 4, 2), "single-primary materialization dims"); + expect(!materialization.ordered_cold_union, "single-primary path keeps existing cold order"); + expect(materialization.cold_expert_ids_by_layer.empty(), "no cold override"); +} + +static void test_builds_shared_expert_split_state_components() { + ExpertSplitConfig cfg; + cfg.n_layer = 2; + cfg.n_expert = 3; + + const std::vector targets = { + {"cuda:0", "cuda", 0, 3, 0, false}, + {"hip:0", "hip", 0, 1, 0, false}, + }; + const std::vector units = { + {0, 0, 1, 100.0}, + {0, 1, 1, 90.0}, + {0, 2, 1, 80.0}, + {1, 0, 1, 70.0}, + {1, 1, 1, 60.0}, + {1, 2, 1, 50.0}, + }; + + ExpertSplitStateComponents state; + std::string err; + expect(build_expert_split_state(cfg, targets, units, /*n_expert_used=*/2, + state, &err), + err.c_str()); + expect(state.matches(2, 3, 2), "shared state matches dims"); + expect(state.plan.at(0, 0).target_index == 0, "plan target lookup"); + expect(state.runtime.target_index(1, 2) == 2, "runtime includes cpu fallback"); + expect(state.compute_runtime.local_index(1, 2) == 1, "compute runtime local id"); + expect(state.materialization.ordered_cold_union, "materialization ordered cold union"); +} + +static void test_builds_generic_layer_mapping() { + ExpertSplitLayerMapping mapping; + std::string err; + expect(build_expert_split_layer_mapping( + /*n_total_layer=*/5, {1, 3}, mapping, &err), + err.c_str()); + expect(mapping.n_total_layer == 5, "mapping total layers"); + expect(mapping.split_layer_count() == 2, "mapping split count"); + expect(mapping.physical_layer_by_split_layer[0] == 1, "mapping split0"); + expect(mapping.physical_layer_by_split_layer[1] == 3, "mapping split1"); + expect(mapping.split_layer_for_physical(1) == 0, "mapping reverse 1"); + expect(mapping.split_layer_for_physical(3) == 1, "mapping reverse 3"); + expect(mapping.split_layer_for_physical(0) == -1, "mapping dense lead"); + expect(mapping.is_split_layer(3), "mapping is split layer"); + expect(!mapping.is_split_layer(4), "mapping rejects dense layer"); +} + +static void test_builds_contiguous_layer_mapping() { + ExpertSplitLayerMapping mapping; + std::string err; + expect(build_contiguous_expert_split_layer_mapping( + /*n_total_layer=*/4, /*first_split_layer=*/1, + /*n_split_layer=*/3, mapping, &err), + err.c_str()); + expect(mapping.physical_layer_by_split_layer.size() == 3, "contiguous mapping size"); + expect(mapping.physical_layer_by_split_layer[0] == 1, "contiguous split0"); + expect(mapping.physical_layer_by_split_layer[1] == 2, "contiguous split1"); + expect(mapping.physical_layer_by_split_layer[2] == 3, "contiguous split2"); + expect(mapping.split_layer_for_physical(1) == 0, "contiguous reverse1"); + expect(mapping.split_layer_for_physical(2) == 1, "contiguous reverse2"); + expect(mapping.split_layer_for_physical(3) == 2, "contiguous reverse3"); +} + static void test_runtime_rejects_invalid_plan() { ExpertSplitPlan plan; plan.n_layer = 1; @@ -251,6 +375,11 @@ int main() { test_compute_runtime_materializes_target_views(); test_compute_runtime_preserves_target_local_order(); test_compute_runtime_fingerprint_changes_with_assignment(); + test_materialization_builds_ordered_cold_union(); + test_materialization_single_explicit_gpu_keeps_single_primary_shape(); + test_builds_shared_expert_split_state_components(); + test_builds_generic_layer_mapping(); + test_builds_contiguous_layer_mapping(); test_runtime_rejects_invalid_plan(); std::printf("OK\n"); return 0; diff --git a/server/test/test_expert_split_state.cpp b/server/test/test_expert_split_state.cpp deleted file mode 100644 index 77f2f88cb..000000000 --- a/server/test/test_expert_split_state.cpp +++ /dev/null @@ -1,86 +0,0 @@ -#include "../src/common/expert_split_state.h" - -#include -#include -#include -#include - -using namespace dflash::common; - -static void expect(bool cond, const char * msg) { - if (!cond) { - std::fprintf(stderr, "FAIL: %s\n", msg); - std::exit(1); - } -} - -static void test_builds_shared_expert_split_state_components() { - ExpertSplitConfig cfg; - cfg.n_layer = 2; - cfg.n_expert = 3; - - const std::vector targets = { - {"cuda:0", "cuda", 0, 3, 0, false}, - {"hip:0", "hip", 0, 1, 0, false}, - }; - const std::vector units = { - {0, 0, 1, 100.0}, - {0, 1, 1, 90.0}, - {0, 2, 1, 80.0}, - {1, 0, 1, 70.0}, - {1, 1, 1, 60.0}, - {1, 2, 1, 50.0}, - }; - - ExpertSplitStateComponents state; - std::string err; - expect(build_expert_split_state(cfg, targets, units, /*n_expert_used=*/2, - state, &err), - err.c_str()); - expect(state.matches(2, 3, 2), "shared state matches dims"); - expect(state.plan.at(0, 0).target_index == 0, "plan target lookup"); - expect(state.runtime.target_index(1, 2) == 2, "runtime includes cpu fallback"); - expect(state.compute_runtime.local_index(1, 2) == 1, "compute runtime local id"); - expect(state.materialization.ordered_cold_union, "materialization ordered cold union"); -} - -static void test_builds_generic_layer_mapping() { - ExpertSplitLayerMapping mapping; - std::string err; - expect(build_expert_split_layer_mapping( - /*n_total_layer=*/5, {1, 3}, mapping, &err), - err.c_str()); - expect(mapping.n_total_layer == 5, "mapping total layers"); - expect(mapping.split_layer_count() == 2, "mapping split count"); - expect(mapping.physical_layer_by_split_layer[0] == 1, "mapping split0"); - expect(mapping.physical_layer_by_split_layer[1] == 3, "mapping split1"); - expect(mapping.split_layer_for_physical(1) == 0, "mapping reverse 1"); - expect(mapping.split_layer_for_physical(3) == 1, "mapping reverse 3"); - expect(mapping.split_layer_for_physical(0) == -1, "mapping dense lead"); - expect(mapping.is_split_layer(3), "mapping is split layer"); - expect(!mapping.is_split_layer(4), "mapping rejects dense layer"); -} - -static void test_builds_contiguous_layer_mapping() { - ExpertSplitLayerMapping mapping; - std::string err; - expect(build_contiguous_expert_split_layer_mapping( - /*n_total_layer=*/4, /*first_split_layer=*/1, - /*n_split_layer=*/3, mapping, &err), - err.c_str()); - expect(mapping.physical_layer_by_split_layer.size() == 3, "contiguous mapping size"); - expect(mapping.physical_layer_by_split_layer[0] == 1, "contiguous split0"); - expect(mapping.physical_layer_by_split_layer[1] == 2, "contiguous split1"); - expect(mapping.physical_layer_by_split_layer[2] == 3, "contiguous split2"); - expect(mapping.split_layer_for_physical(1) == 0, "contiguous reverse1"); - expect(mapping.split_layer_for_physical(2) == 1, "contiguous reverse2"); - expect(mapping.split_layer_for_physical(3) == 2, "contiguous reverse3"); -} - -int main() { - test_builds_shared_expert_split_state_components(); - test_builds_generic_layer_mapping(); - test_builds_contiguous_layer_mapping(); - std::printf("OK\n"); - return 0; -}