From 66ab05cfe2c9253a603e934c681bedb8050ac68b Mon Sep 17 00:00:00 2001 From: weicj Date: Sat, 23 May 2026 19:43:58 +0800 Subject: [PATCH 1/7] feat(server): add target-layer-split backend adapter path --- server/CMakeLists.txt | 2 + server/src/common/backend_factory.cpp | 27 ++ server/src/common/dflash_spec_decode.cpp | 40 +- server/src/common/dflash_spec_decode.h | 19 +- server/src/common/layer_split_backend.cpp | 225 ++++++++++ server/src/common/layer_split_backend.h | 112 +++++ server/src/common/layer_split_utils.cpp | 24 +- server/src/common/layer_split_utils.h | 4 +- server/src/qwen35/layer_split_forward.cpp | 4 + .../src/qwen35/qwen35_layer_split_adapter.cpp | 418 ++++++++++++++++++ .../src/qwen35/qwen35_layer_split_adapter.h | 105 +++++ server/src/qwen35/qwen35_target_graph.cpp | 29 +- server/src/server/server_main.cpp | 25 +- server/test/test_server_unit.cpp | 202 +++++++++ 14 files changed, 1211 insertions(+), 25 deletions(-) create mode 100644 server/src/common/layer_split_backend.cpp create mode 100644 server/src/common/layer_split_backend.h create mode 100644 server/src/qwen35/qwen35_layer_split_adapter.cpp create mode 100644 server/src/qwen35/qwen35_layer_split_adapter.h diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 8ce0344dd..4be4744a4 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -244,6 +244,7 @@ add_library(dflash_common STATIC src/common/pflash_drafter_ipc.cpp src/common/dflash_draft_graph.cpp src/common/dflash_spec_decode.cpp + src/common/layer_split_backend.cpp src/qwen35/graph_builders.cpp src/qwen35moe/qwen35moe_ffn.cpp src/qwen35moe/qwen35moe_backend.cpp @@ -256,6 +257,7 @@ add_library(dflash_common STATIC src/qwen35/layer_split_forward.cpp src/qwen35/layer_split_daemon.cpp src/qwen35/qwen35_backend.cpp + src/qwen35/qwen35_layer_split_adapter.cpp src/qwen35/qwen35_dflash_target.cpp src/qwen35/qwen35_layer_split_dflash_target.cpp src/qwen35/layer_split_daemon_loop.cpp diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 0fea6cc37..2075fd1cc 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -8,8 +8,11 @@ #include "laguna_backend.h" #include "qwen3_backend.h" #include "gemma4_backend.h" +#include "layer_split_backend.h" +#include "qwen35_layer_split_adapter.h" #include +#include namespace dflash::common { @@ -42,6 +45,30 @@ std::unique_ptr create_backend(const BackendArgs & args) { std::fprintf(stderr, "[backend_factory] detected arch=%s\n", arch.c_str()); if (arch == "qwen35") { + if (args.device.is_layer_split()) { + Qwen35LayerSplitAdapterConfig cfg; + cfg.target_path = args.model_path; + cfg.draft_path = args.draft_path; + cfg.device = args.device; + cfg.draft_gpu = args.draft_device.gpu; + cfg.remote_draft = args.remote_draft; + cfg.fa_window = args.fa_window; + cfg.kq_stride_pad = args.kq_stride_pad; + cfg.draft_ctx_max = args.draft_ctx_max; + cfg.max_verify_tokens = args.ddtree_mode + ? std::max(DFLASH27B_DRAFT_BLOCK_SIZE, args.ddtree_budget + 1) + : DFLASH27B_DRAFT_BLOCK_SIZE; + cfg.run_dflash = args.draft_path != nullptr; + + auto adapter = std::make_unique(cfg); + auto backend = std::make_unique(std::move(adapter)); + if (!backend->init()) { + std::fprintf(stderr, "[backend_factory] LayerSplitBackend(qwen35) init failed\n"); + return nullptr; + } + return backend; + } + Qwen35Config cfg; cfg.target_path = args.model_path; cfg.draft_path = args.draft_path; diff --git a/server/src/common/dflash_spec_decode.cpp b/server/src/common/dflash_spec_decode.cpp index 55d2a3556..141e45e92 100644 --- a/server/src/common/dflash_spec_decode.cpp +++ b/server/src/common/dflash_spec_decode.cpp @@ -34,7 +34,30 @@ bool run_dflash_spec_decode( int draft_ctx_max, int stream_fd, DFlashDraftIpcClient * remote_draft, - const std::vector * hint_tokens) { + const std::vector * hint_tokens, + int base_pos) { + DaemonIO io; + io.stream_fd = stream_fd; + return run_dflash_spec_decode(target, draft_weights, draft_backend, + feature_ring, prompt, n_gen, last_tok, + out_path, draft_ctx_max, io, + remote_draft, hint_tokens, base_pos); +} + +bool run_dflash_spec_decode( + DFlashTarget & target, + DraftWeights & draft_weights, + ggml_backend_t draft_backend, + DraftFeatureMirror & feature_ring, + const std::vector & prompt, + int n_gen, + int last_tok, + const char * out_path, + int draft_ctx_max, + const DaemonIO & io, + DFlashDraftIpcClient * remote_draft, + const std::vector * hint_tokens, + int base_pos) { const bool use_remote_draft = remote_draft && remote_draft->active(); if (!use_remote_draft && !feature_ring.target_feat) return false; @@ -54,7 +77,7 @@ bool run_dflash_spec_decode( std::vector remote_hidden; // host buffer for remote-draft hidden states std::vector out_all = prompt; - int committed = (int)prompt.size(); + int committed = base_pos + (int)prompt.size(); int n_generated = 0; int n_draft_steps = 0; int n_accept_sum = 0; @@ -199,15 +222,19 @@ bool run_dflash_spec_decode( last_tok = replay_last_tok; bool hit_eos = false; + int emitted = 0; for (int i = 0; i < commit_n; i++) { out_all.push_back(replay_tok[i]); - stream_emit_fd(stream_fd, replay_tok[i]); + io.emit(replay_tok[i]); + if (io.cancelled) break; + ++emitted; if (target.is_eos(replay_tok[i])) hit_eos = true; } - committed += commit_n; - n_generated += commit_n; - n_accept_sum += std::min(accept_n, commit_n); + committed += emitted; + n_generated += emitted; + n_accept_sum += std::min(accept_n, emitted); n_draft_steps++; + if (io.cancelled) break; if (hit_eos) break; } if (!use_remote_draft && draft_backend) ggml_backend_synchronize(draft_backend); @@ -231,4 +258,3 @@ bool run_dflash_spec_decode( } } // namespace dflash::common - diff --git a/server/src/common/dflash_spec_decode.h b/server/src/common/dflash_spec_decode.h index 15dad3fdf..6546010d8 100644 --- a/server/src/common/dflash_spec_decode.h +++ b/server/src/common/dflash_spec_decode.h @@ -14,6 +14,7 @@ #include "dflash_target.h" #include "dflash_feature_ring.h" #include "dflash_draft_ipc.h" +#include "model_backend.h" #include "ggml.h" #include "ggml-backend.h" @@ -53,6 +54,22 @@ bool run_dflash_spec_decode( int draft_ctx_max, int stream_fd = -1, DFlashDraftIpcClient * remote_draft = nullptr, - const std::vector * hint_tokens = nullptr); + const std::vector * hint_tokens = nullptr, + int base_pos = 0); + +bool run_dflash_spec_decode( + DFlashTarget & target, + DraftWeights & draft_weights, + ggml_backend_t draft_backend, + DraftFeatureMirror & feature_ring, + const std::vector & prompt, + int n_gen, + int last_tok, + const char * out_path, + int draft_ctx_max, + const DaemonIO & io, + DFlashDraftIpcClient * remote_draft = nullptr, + const std::vector * hint_tokens = nullptr, + int base_pos = 0); } // namespace dflash::common diff --git a/server/src/common/layer_split_backend.cpp b/server/src/common/layer_split_backend.cpp new file mode 100644 index 000000000..13cbc2c87 --- /dev/null +++ b/server/src/common/layer_split_backend.cpp @@ -0,0 +1,225 @@ +// Generic server-facing backend for target layer split. + +#include "layer_split_backend.h" + +#include "io_utils.h" + +#include +#include +#include + +namespace dflash::common { + +LayerSplitBackend::LayerSplitBackend(std::unique_ptr adapter) + : adapter_(std::move(adapter)) {} + +LayerSplitBackend::~LayerSplitBackend() { shutdown(); } + +bool LayerSplitBackend::init() { + if (!adapter_) { + std::fprintf(stderr, "[target-split] missing model adapter\n"); + return false; + } + return adapter_->init(); +} + +void LayerSplitBackend::print_ready_banner() const { + std::printf("[daemon] ready\n"); + std::fflush(stdout); +} + +bool LayerSplitBackend::park(const std::string & what) { + std::fprintf(stderr, "[target-split] park is not supported yet (%s)\n", + what.c_str()); + return false; +} + +bool LayerSplitBackend::unpark(const std::string & what) { + std::fprintf(stderr, "[target-split] unpark is not supported yet (%s)\n", + what.c_str()); + return false; +} + +GenerateResult LayerSplitBackend::run_from_state(const GenerateRequest & req, + const DaemonIO & io, + int base_pos, + bool reset_state) { + GenerateResult result; + if (!adapter_) { + result.error = "adapter"; + return result; + } + + DaemonIO out_io = io.with_token_callback(req.on_token); + if (base_pos + (int)req.prompt.size() + req.n_gen + 1 > adapter_->max_context()) { + result.error = "context"; + return result; + } + if (req.do_sample && req.sampler.temp > 0.0f) { + result.error = "sampling_unsupported"; + return result; + } + + adapter_->begin_request(req); + if (reset_state) adapter_->reset_request_state(); + + const int prompt_len = (int)req.prompt.size(); + int last_tok = (base_pos > 0 && prompt_len == 0) + ? adapter_->current_last_token() + : -1; + int consumed = 0; + auto t_prefill_start = std::chrono::steady_clock::now(); + while (consumed < prompt_len) { + int n_tokens = prompt_len - consumed; + if (req.snap_pos >= 0 && req.snap_slot >= 0 && + req.snap_pos > base_pos + consumed && + req.snap_pos < base_pos + consumed + n_tokens) { + n_tokens = req.snap_pos - (base_pos + consumed); + } + std::vector chunk(req.prompt.begin() + consumed, + req.prompt.begin() + consumed + n_tokens); + if (!adapter_->prefill(chunk, base_pos + consumed, last_tok)) { + result.error = "prefill"; + return result; + } + consumed += n_tokens; + if (req.snap_pos >= 0 && req.snap_slot >= 0 && + base_pos + consumed == req.snap_pos) { + if (adapter_->snapshot_save(req.snap_slot)) { + std::printf("[snap] inline slot=%d cur_pos=%d\n", + req.snap_slot, req.snap_pos); + std::fflush(stdout); + } + } + } + result.prefill_s = std::chrono::duration( + std::chrono::steady_clock::now() - t_prefill_start).count(); + + if (req.n_gen > 0) { + if (last_tok < 0) { + result.error = "decode_seed"; + return result; + } + auto t_decode_start = std::chrono::steady_clock::now(); + const bool ok = (base_pos == 0 && adapter_->can_dflash_decode()) + ? adapter_->decode_dflash(req.prompt, base_pos, last_tok, req.n_gen, + result.tokens, out_io) + : adapter_->decode_ar(last_tok, base_pos + (int)req.prompt.size(), req.n_gen, + result.tokens, out_io); + if (!ok) { + result.error = "decode"; + return result; + } + result.decode_s = std::chrono::duration( + std::chrono::steady_clock::now() - t_decode_start).count(); + } + + result.ok = true; + return result; +} + +GenerateResult LayerSplitBackend::generate(const GenerateRequest & req, + const DaemonIO & io) { + return run_from_state(req, io, /*base_pos=*/0, /*reset_state=*/true); +} + +bool LayerSplitBackend::snapshot_save(int slot) { + return adapter_ && adapter_->snapshot_save(slot); +} + +void LayerSplitBackend::snapshot_free(int slot) { + if (adapter_) adapter_->snapshot_free(slot); +} + +bool LayerSplitBackend::snapshot_used(int slot) const { + return adapter_ && adapter_->snapshot_used(slot); +} + +int LayerSplitBackend::snapshot_cur_pos(int slot) const { + return adapter_ ? adapter_->snapshot_cur_pos(slot) : 0; +} + +GenerateResult LayerSplitBackend::restore_and_generate( + int slot, const GenerateRequest & req, const DaemonIO & io) { + GenerateResult result; + if (!adapter_ || !adapter_->snapshot_restore(slot)) { + result.error = "bad slot"; + io.emit(-1); + return result; + } + const int snap_pos = adapter_->snapshot_cur_pos(slot); + if ((int)req.prompt.size() < snap_pos) { + result.error = "snapshot_longer_than_prompt"; + io.emit(-1); + return result; + } + GenerateRequest delta_req = req; + delta_req.prompt = std::vector( + req.prompt.begin() + snap_pos, req.prompt.end()); + return run_from_state(delta_req, io, snap_pos, /*reset_state=*/false); +} + +ModelBackend::CompressResult +LayerSplitBackend::compress(const CompressRequest & req) { + return adapter_ ? adapter_->compress(req) : CompressResult{}; +} + +bool LayerSplitBackend::handle_compress(const std::string & line, + const DaemonIO & io) { + std::string args = line.size() > 9 ? line.substr(9) : std::string{}; + bool skip_park = false; + const std::string suffix = " nopark"; + if (args.size() >= suffix.size() && + args.compare(args.size() - suffix.size(), suffix.size(), suffix) == 0) { + skip_park = true; + args.resize(args.size() - suffix.size()); + } + + char ppath[1024]; + int keep_x1000 = 0; + char drafter_path[1024] = {0}; + const int n = std::sscanf(args.c_str(), "%1023s %d %1023s", + ppath, &keep_x1000, drafter_path); + if (n < 2) { + std::fprintf(stderr, "[target-split][compress] bad args\n"); + io.emit(-1); + return false; + } + + CompressRequest req; + req.input_ids = read_int32_file(ppath); + req.keep_ratio = (float)keep_x1000 / 1000.0f; + if (n >= 3 && drafter_path[0]) { + req.drafter_path = drafter_path; + } else if (adapter_) { + req.drafter_path = adapter_->default_compress_drafter_path(); + } + req.skip_park = skip_park; + + CompressResult result = compress(req); + for (int32_t t : result.compressed_ids) io.emit(t); + io.emit(-1); + return result.ok; +} + +void LayerSplitBackend::free_drafter() { + if (adapter_) adapter_->free_drafter(); +} + +bool LayerSplitBackend::supports_dflash_spec_decode() const { + return adapter_ && adapter_->supports_dflash_spec_decode(); +} + +DFlashTarget * LayerSplitBackend::dflash_target() { + return adapter_ ? adapter_->dflash_target() : nullptr; +} + +bool LayerSplitBackend::supports_remote_draft() const { + return adapter_ && adapter_->supports_remote_draft(); +} + +void LayerSplitBackend::shutdown() { + if (adapter_) adapter_->shutdown(); +} + +} // namespace dflash::common diff --git a/server/src/common/layer_split_backend.h b/server/src/common/layer_split_backend.h new file mode 100644 index 000000000..4553b42e3 --- /dev/null +++ b/server/src/common/layer_split_backend.h @@ -0,0 +1,112 @@ +// Generic server-facing backend for target layer split. +// +// Model-specific layer-split details live behind LayerSplitAdapter. This keeps +// server placement, request flow, and compatibility policy in one place while +// allowing each architecture to provide only its partial-load/forward/cache +// implementation. + +#pragma once + +#include "model_backend.h" + +#include +#include +#include + +namespace dflash::common { + +class LayerSplitAdapter { +public: + virtual ~LayerSplitAdapter() = default; + + virtual const char * name() const = 0; + virtual bool init() = 0; + virtual int max_context() const = 0; + + virtual void begin_request(const GenerateRequest & req) { (void)req; } + virtual void reset_request_state() = 0; + virtual bool prefill(const std::vector & prompt, + int base_pos, int & last_tok) = 0; + virtual bool decode_ar(int last_tok, int committed, int n_gen, + std::vector & out_tokens, + const DaemonIO & io) = 0; + + virtual bool can_dflash_decode() const { return false; } + virtual bool decode_dflash(const std::vector & prompt, + int base_pos, int last_tok, int n_gen, + std::vector & out_tokens, + const DaemonIO & io) { + (void)prompt; (void)base_pos; (void)last_tok; (void)n_gen; + (void)out_tokens; (void)io; + return false; + } + + virtual bool supports_dflash_spec_decode() const { return false; } + virtual DFlashTarget * dflash_target() { return nullptr; } + virtual bool supports_remote_draft() const { return false; } + + virtual const char * default_compress_drafter_path() const { return ""; } + virtual ModelBackend::CompressResult + compress(const ModelBackend::CompressRequest & req) { + (void)req; + return {}; + } + virtual void free_drafter() = 0; + + virtual bool snapshot_save(int slot) { (void)slot; return false; } + virtual void snapshot_free(int slot) { (void)slot; } + virtual bool snapshot_used(int slot) const { (void)slot; return false; } + virtual int snapshot_cur_pos(int slot) const { (void)slot; return 0; } + virtual bool snapshot_restore(int slot) { (void)slot; return false; } + virtual int current_last_token() const { return -1; } + + virtual void shutdown() = 0; +}; + +class LayerSplitBackend : public ModelBackend { +public: + explicit LayerSplitBackend(std::unique_ptr adapter); + ~LayerSplitBackend() override; + + LayerSplitBackend(const LayerSplitBackend &) = delete; + LayerSplitBackend & operator=(const LayerSplitBackend &) = delete; + + bool init(); + + void print_ready_banner() const override; + bool park(const std::string & what) override; + bool unpark(const std::string & what) override; + bool is_target_parked() const override { return false; } + + GenerateResult generate(const GenerateRequest & req, + const DaemonIO & io) override; + + bool snapshot_save(int slot) override; + void snapshot_free(int slot) override; + bool snapshot_used(int slot) const override; + int snapshot_cur_pos(int slot) const override; + GenerateResult restore_and_generate(int slot, + const GenerateRequest & req, + const DaemonIO & io) override; + + CompressResult compress(const CompressRequest & req) override; + bool handle_compress(const std::string & line, + const DaemonIO & io) override; + void free_drafter() override; + + bool supports_dflash_spec_decode() const override; + DFlashTarget * dflash_target() override; + bool supports_remote_draft() const override; + + void shutdown() override; + +private: + GenerateResult run_from_state(const GenerateRequest & req, + const DaemonIO & io, + int base_pos, + bool reset_state); + + std::unique_ptr adapter_; +}; + +} // namespace dflash::common diff --git a/server/src/common/layer_split_utils.cpp b/server/src/common/layer_split_utils.cpp index 57bf0db49..f69501bf0 100644 --- a/server/src/common/layer_split_utils.cpp +++ b/server/src/common/layer_split_utils.cpp @@ -41,13 +41,19 @@ std::vector> compute_layer_ranges( std::string validate_device_placement( const DevicePlacement & dp, - int cuda_device_count) + int device_count) { - if (cuda_device_count <= 0) return "no CUDA devices available"; + const bool validate_device_count = device_count >= 0; + if (validate_device_count && device_count == 0) { + return "no GPU devices available"; + } - if (dp.gpu < 0 || dp.gpu >= cuda_device_count) { - return "primary gpu " + std::to_string(dp.gpu) + - " out of range [0, " + std::to_string(cuda_device_count) + ")"; + if (dp.gpu < 0 || + (validate_device_count && dp.gpu >= device_count)) { + return "primary gpu " + std::to_string(dp.gpu) + " out of range" + + (validate_device_count + ? " [0, " + std::to_string(device_count) + ")" + : ""); } if (!dp.layer_split_gpus.empty()) { @@ -57,9 +63,13 @@ std::string validate_device_placement( std::set seen; for (int g : dp.layer_split_gpus) { - if (g < 0 || g >= cuda_device_count) { + if (g < 0 || + (validate_device_count && g >= device_count)) { return "layer_split gpu " + std::to_string(g) + - " out of range [0, " + std::to_string(cuda_device_count) + ")"; + " out of range" + + (validate_device_count + ? " [0, " + std::to_string(device_count) + ")" + : ""); } if (!seen.insert(g).second) { return "duplicate gpu " + std::to_string(g) + " in layer_split_gpus"; diff --git a/server/src/common/layer_split_utils.h b/server/src/common/layer_split_utils.h index 1081b65ab..98ff422af 100644 --- a/server/src/common/layer_split_utils.h +++ b/server/src/common/layer_split_utils.h @@ -23,9 +23,11 @@ std::vector> compute_layer_ranges( const std::vector & weights); // Validate a DevicePlacement against system constraints. +// If device_count is negative, only validates structural constraints that do +// not require querying the runtime-visible GPU count. // Returns empty string on success, error description on failure. std::string validate_device_placement( const DevicePlacement & dp, - int cuda_device_count); + int device_count); } // namespace dflash::common diff --git a/server/src/qwen35/layer_split_forward.cpp b/server/src/qwen35/layer_split_forward.cpp index 985be9502..1ee664301 100644 --- a/server/src/qwen35/layer_split_forward.cpp +++ b/server/src/qwen35/layer_split_forward.cpp @@ -215,6 +215,10 @@ bool run_target_layer_split_forward( activation_pair_free(acts); if (!ok) return false; last_tok = argmax_tokens.empty() ? -1 : argmax_tokens.back(); + for (auto & shard : shards) { + shard.cache.cur_pos = base_pos + n_tokens_total; + shard.cache.last_tok = last_tok; + } if (argmax_out) *argmax_out = std::move(argmax_tokens); if (logits_out) logits_out->clear(); return true; diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp new file mode 100644 index 000000000..0d6dcc3d0 --- /dev/null +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -0,0 +1,418 @@ +// Qwen35 layer-split adapter. + +#include "qwen35_layer_split_adapter.h" + +#include "common/dflash_spec_decode.h" +#include "common/gguf_inspect.h" +#include "common/layer_split_utils.h" +#include "common/peer_access.h" +#include "common/sampler.h" +#include "common/snapshot_backend.h" +#include "qwen35/layer_split_forward.h" +#include "qwen35/qwen35_layer_split_dflash_target.h" +#include "qwen3/qwen3_drafter.h" + +#include "ggml-cuda.h" + +#include +#include +#include + +namespace dflash::common { + +Qwen35LayerSplitAdapter::Qwen35LayerSplitAdapter( + const Qwen35LayerSplitAdapterConfig & cfg) + : cfg_(cfg) {} + +Qwen35LayerSplitAdapter::~Qwen35LayerSplitAdapter() { shutdown(); } + +bool Qwen35LayerSplitAdapter::init() { + if (!cfg_.target_path || cfg_.device.layer_split_gpus.size() < 2) { + std::fprintf(stderr, "[target-split] invalid layer-split config\n"); + return false; + } + + const auto info = inspect_gguf_model_info(cfg_.target_path); + const int n_layer = info.n_layer; + if (n_layer <= 0) { + std::fprintf(stderr, "[target-split] failed to inspect target layer count\n"); + return false; + } + const auto ranges = compute_layer_ranges( + n_layer, + (int)cfg_.device.layer_split_gpus.size(), + cfg_.device.layer_split_weights); + if (ranges.size() != cfg_.device.layer_split_gpus.size()) { + std::fprintf(stderr, + "[target-split] bad layer split for %zu GPUs and %d layers\n", + cfg_.device.layer_split_gpus.size(), n_layer); + return false; + } + + shards_.resize(cfg_.device.layer_split_gpus.size()); + for (size_t i = 0; i < shards_.size(); ++i) { + auto & shard = shards_[i]; + shard.gpu = cfg_.device.layer_split_gpus[i]; + shard.layer_begin = ranges[i].first; + shard.layer_end = ranges[i].second; + shard.backend = ggml_backend_cuda_init(shard.gpu); + if (!shard.backend) { + std::fprintf(stderr, + "[target-split] backend init failed gpu=%d\n", shard.gpu); + return false; + } + } + + if (cfg_.device.peer_access) { + for (size_t i = 0; i < cfg_.device.layer_split_gpus.size(); ++i) { + for (size_t j = i + 1; j < cfg_.device.layer_split_gpus.size(); ++j) { + (void)enable_peer_access_pair(cfg_.device.layer_split_gpus[i], + cfg_.device.layer_split_gpus[j]); + } + } + } + + snapshot_backends_.resize(shards_.size(), nullptr); + for (size_t i = 0; i < shards_.size(); ++i) { + snapshot_backends_[i] = create_snapshot_backend(shards_[i].backend); + if (!snapshot_backends_[i]) { + std::fprintf(stderr, + "[target-split] snapshot backend init failed gpu=%d\n", + shards_[i].gpu); + return false; + } + } + + for (auto & shard : shards_) { + TargetLoadPlan plan; + plan.layer_begin = shard.layer_begin; + plan.layer_end = shard.layer_end; + plan.load_output = (&shard == &shards_.back()); + if (!load_target_gguf_partial(cfg_.target_path, shard.backend, plan, + shard.weights) || + !create_target_cache_partial(shard.weights, cfg_.device.max_ctx, + cfg_.max_verify_tokens, shard.backend, + shard.cache, + /*prefill_only=*/!cfg_.run_dflash, + shard.layer_begin, shard.layer_end, + /*allocate_target_feat=*/false)) { + std::fprintf(stderr, "[target-split] load/cache gpu=%d: %s\n", + shard.gpu, dflash27b_last_error()); + return false; + } + std::fprintf(stderr, "[target-split] gpu=%d layers=[%d,%d)\n", + shard.gpu, shard.layer_begin, shard.layer_end); + } + + if (cfg_.draft_path && cfg_.run_dflash && !load_draft()) { + return false; + } + prefix_snapshots_.resize(PREFIX_SLOTS); + for (auto & slot : prefix_snapshots_) { + slot.resize(shards_.size()); + } + + return true; +} + +bool Qwen35LayerSplitAdapter::load_draft() { + if (cfg_.remote_draft.enabled()) { + const int cap = cfg_.remote_draft.ring_cap > 0 + ? std::min(cfg_.remote_draft.ring_cap, cfg_.device.max_ctx) + : std::min(cfg_.device.max_ctx, cfg_.draft_ctx_max); + if (!remote_draft_.start(cfg_.remote_draft.ipc_bin, cfg_.draft_path, + cfg_.draft_gpu, cap, + cfg_.remote_draft.work_dir)) { + std::fprintf(stderr, + "[target-split] remote draft start failed gpu=%d\n", + cfg_.draft_gpu); + return false; + } + draft_weights_.n_embd = DFLASH27B_TARGET_HIDDEN; + draft_weights_.block_size = DFLASH27B_DRAFT_BLOCK_SIZE; + draft_weights_.n_target_layers = DFLASH27B_DRAFT_N_TARGET_LAYERS; + std::fprintf(stderr, + "[target-split] remote draft ready gpu=%d cap=%d\n", + cfg_.draft_gpu, cap); + return true; + } + + for (auto & shard : shards_) { + if (shard.gpu == cfg_.draft_gpu) { + draft_backend_ = shard.backend; + break; + } + } + if (!draft_backend_) { + draft_backend_ = ggml_backend_cuda_init(cfg_.draft_gpu); + if (!draft_backend_) { + std::fprintf(stderr, + "[target-split] draft backend init failed gpu=%d\n", + cfg_.draft_gpu); + return false; + } + draft_backend_owned_ = true; + } + + std::string draft_path(cfg_.draft_path ? cfg_.draft_path : ""); + const bool draft_ok = draft_path.size() >= 5 && + draft_path.substr(draft_path.size() - 5) == ".gguf" + ? load_draft_gguf(cfg_.draft_path, draft_backend_, draft_weights_, + &shards_.front().weights) + : load_draft_safetensors(cfg_.draft_path, draft_backend_, + draft_weights_, &shards_.front().weights); + if (!draft_ok) { + std::fprintf(stderr, "[target-split] draft load gpu=%d: %s\n", + cfg_.draft_gpu, dflash27b_last_error()); + return false; + } + + const int cap = std::min(cfg_.device.max_ctx, cfg_.draft_ctx_max); + if (!draft_feature_mirror_init(feature_ring_, draft_backend_, + cfg_.draft_gpu, cfg_.draft_gpu, cap, + draft_weights_.n_target_layers, + draft_weights_.n_embd)) { + std::fprintf(stderr, + "[target-split] draft feature ring init failed gpu=%d\n", + cfg_.draft_gpu); + return false; + } + return true; +} + +void Qwen35LayerSplitAdapter::begin_request(const GenerateRequest & req) { + sampler_ = req.sampler; + if (req.do_sample && sampler_.seed != 0) { + sampler_rng_.seed(sampler_.seed); + } +} + +void Qwen35LayerSplitAdapter::reset_request_state() { + for (auto & shard : shards_) reset_target_cache(shard.cache); +} + +bool Qwen35LayerSplitAdapter::prefill(const std::vector & prompt, + int base_pos, int & last_tok) { + if (prompt.empty()) return false; + if (base_pos < 0 || base_pos + (int)prompt.size() > cfg_.device.max_ctx) { + std::fprintf(stderr, + "[target-split] prompt range [%d,%zu) exceeds max_ctx (%d)\n", + base_pos, (size_t)base_pos + prompt.size(), cfg_.device.max_ctx); + return false; + } + int ubatch = prompt.size() > 2048 ? 384 : 16; + if (const char * s = std::getenv("DFLASH27B_PREFILL_UBATCH")) { + ubatch = std::max(1, std::atoi(s)); + } + return run_target_layer_split_forward( + shards_, shards_.front().weights, prompt, base_pos, ubatch, last_tok, + cfg_.kq_stride_pad, /*fa_window=*/0, + (cfg_.run_dflash && !remote_draft_.active()) ? &feature_ring_ : nullptr, + /*argmax_out=*/nullptr, /*logits_out=*/nullptr, + cfg_.run_dflash ? &remote_draft_ : nullptr); +} + +bool Qwen35LayerSplitAdapter::snapshot_slot_valid(int slot) const { + return slot >= 0 && slot < PREFIX_SLOTS && + prefix_snapshots_.size() == (size_t)PREFIX_SLOTS && + !shards_.empty(); +} + +bool Qwen35LayerSplitAdapter::snapshot_save(int slot) { + if (!snapshot_slot_valid(slot)) return false; + if (snapshot_backends_.size() != shards_.size()) return false; + snapshot_free(slot); + auto & snaps = prefix_snapshots_[(size_t)slot]; + if (snaps.size() != shards_.size()) snaps.resize(shards_.size()); + for (size_t i = 0; i < shards_.size(); ++i) { + if (!snapshot_target_cache(shards_[i].weights, shards_[i].cache, + snapshot_backends_[i], snaps[i])) { + for (size_t j = 0; j <= i && j < snaps.size(); ++j) { + free_prefix_snapshot(snaps[j]); + } + return false; + } + } + return true; +} + +void Qwen35LayerSplitAdapter::snapshot_free(int slot) { + if (!snapshot_slot_valid(slot)) return; + for (auto & snap : prefix_snapshots_[(size_t)slot]) { + free_prefix_snapshot(snap); + } +} + +bool Qwen35LayerSplitAdapter::snapshot_used(int slot) const { + if (!snapshot_slot_valid(slot)) return false; + const auto & snaps = prefix_snapshots_[(size_t)slot]; + if (snaps.size() != shards_.size()) return false; + for (const auto & snap : snaps) { + if (!snap.ctx) return false; + } + return true; +} + +int Qwen35LayerSplitAdapter::snapshot_cur_pos(int slot) const { + if (!snapshot_used(slot)) return 0; + return prefix_snapshots_[(size_t)slot].front().cur_pos; +} + +bool Qwen35LayerSplitAdapter::snapshot_restore(int slot) { + if (!snapshot_used(slot)) return false; + auto & snaps = prefix_snapshots_[(size_t)slot]; + const int cur_pos = snaps.front().cur_pos; + for (size_t i = 0; i < shards_.size(); ++i) { + if (snaps[i].cur_pos != cur_pos || + !restore_target_cache(snaps[i], shards_[i].cache)) { + return false; + } + } + return true; +} + +int Qwen35LayerSplitAdapter::current_last_token() const { + if (shards_.empty()) return -1; + return shards_.front().cache.last_tok; +} + +bool Qwen35LayerSplitAdapter::decode_ar( + int last_tok, int committed, int n_gen, + std::vector & out_tokens, + const DaemonIO & io) { + if (n_gen <= 0) return true; + + out_tokens.push_back(last_tok); + io.emit(last_tok); + if (io.cancelled) { + io.emit(-1); + return true; + } + if (is_eos_tok(last_tok, shards_.front().weights)) { + io.emit(-1); + return true; + } + ++committed; + + for (int i = 1; i < n_gen; ++i) { + std::vector one(1, last_tok); + int next_tok = -1; + if (!run_target_layer_split_forward( + shards_, shards_.front().weights, one, committed, 1, next_tok, + cfg_.kq_stride_pad, cfg_.fa_window, + cfg_.run_dflash ? &feature_ring_ : nullptr)) { + return false; + } + out_tokens.push_back(next_tok); + io.emit(next_tok); + if (io.cancelled) break; + if (is_eos_tok(next_tok, shards_.front().weights)) break; + last_tok = next_tok; + ++committed; + } + io.emit(-1); + return true; +} + +bool Qwen35LayerSplitAdapter::can_dflash_decode() const { + return cfg_.run_dflash && cfg_.draft_path && sampler_.temp == 0.0f; +} + +bool Qwen35LayerSplitAdapter::decode_dflash( + const std::vector & prompt, int base_pos, int last_tok, int n_gen, + std::vector & out_tokens, const DaemonIO & io) { + const bool use_remote_draft = remote_draft_.active(); + Qwen35LayerSplitDFlashTarget target( + shards_, use_remote_draft ? nullptr : &feature_ring_, + cfg_.kq_stride_pad, cfg_.fa_window, + use_remote_draft ? &remote_draft_ : nullptr); + DaemonIO collect_io = io.with_token_callback([&](int32_t tok) -> bool { + out_tokens.push_back(tok); + return true; + }); + const bool ok = run_dflash_spec_decode( + target, draft_weights_, draft_backend_, feature_ring_, prompt, n_gen, + last_tok, /*out_path=*/nullptr, cfg_.draft_ctx_max, collect_io, + use_remote_draft ? &remote_draft_ : nullptr, /*hint_tokens=*/nullptr, base_pos); + return ok; +} + +const char * Qwen35LayerSplitAdapter::default_compress_drafter_path() const { + return "/opt/lucebox/models/drafter/Qwen3-0.6B-BF16.gguf"; +} + +ModelBackend::CompressResult +Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) { + ModelBackend::CompressResult result; + if (req.input_ids.empty() || req.drafter_path.empty()) return result; + + for (auto & shard : shards_) ggml_backend_synchronize(shard.backend); + if (draft_backend_) ggml_backend_synchronize(draft_backend_); + + if (!pflash_drafter_loaded_) { + std::fprintf(stderr, "[target-split][compress] loading drafter from %s ...\n", + req.drafter_path.c_str()); + if (!load_drafter(req.drafter_path, /*gpu_layers=*/999, + pflash_drafter_)) { + std::fprintf(stderr, + "[target-split][compress] drafter init failed: %s\n", + dflash27b_last_error()); + return result; + } + pflash_drafter_loaded_ = true; + std::fprintf(stderr, "[target-split][compress] drafter ready\n"); + } + + result.compressed_ids = drafter_score_and_compress( + pflash_drafter_, req.input_ids, req.keep_ratio); + result.ok = !result.compressed_ids.empty(); + if (result.ok) { + std::fprintf(stderr, "[target-split][compress] %zu -> %zu tokens\n", + req.input_ids.size(), result.compressed_ids.size()); + } + return result; +} + +void Qwen35LayerSplitAdapter::free_drafter() { + remote_draft_.close(); + if (pflash_drafter_loaded_) { + dflash::common::free_drafter(pflash_drafter_); + pflash_drafter_loaded_ = false; + } + step_graph_destroy(draft_sg_); + step_graph_destroy(proj_sg_); +} + +DFlashTarget * Qwen35LayerSplitAdapter::dflash_target() { + if (!dflash_target_) { + dflash_target_ = std::make_unique( + shards_, + (cfg_.run_dflash && !remote_draft_.active()) ? &feature_ring_ : nullptr, + cfg_.kq_stride_pad, cfg_.fa_window, + remote_draft_.active() ? &remote_draft_ : nullptr); + } + return dflash_target_.get(); +} + +void Qwen35LayerSplitAdapter::shutdown() { + dflash_target_.reset(); + free_drafter(); + draft_feature_mirror_free(feature_ring_); + free_draft_weights(draft_weights_); + for (auto & slot : prefix_snapshots_) { + for (auto & snap : slot) free_prefix_snapshot(snap); + } + prefix_snapshots_.clear(); + for (size_t i = 0; i < snapshot_backends_.size() && i < shards_.size(); ++i) { + free_snapshot_backend(snapshot_backends_[i], shards_[i].backend); + } + snapshot_backends_.clear(); + if (draft_backend_owned_ && draft_backend_) { + ggml_backend_free(draft_backend_); + } + draft_backend_ = nullptr; + draft_backend_owned_ = false; + free_target_layer_split_shards(shards_); +} + +} // namespace dflash::common diff --git a/server/src/qwen35/qwen35_layer_split_adapter.h b/server/src/qwen35/qwen35_layer_split_adapter.h new file mode 100644 index 000000000..f204c06ae --- /dev/null +++ b/server/src/qwen35/qwen35_layer_split_adapter.h @@ -0,0 +1,105 @@ +// Qwen35 layer-split adapter. + +#pragma once + +#include "common/dflash_draft_ipc.h" +#include "common/layer_split_backend.h" +#include "dflash_feature_ring.h" +#include "layer_split_types.h" +#include "placement/placement_config.h" +#include "placement/remote_draft_config.h" +#include "qwen3/qwen3_drafter.h" +#include "step_graph.h" +#include "internal.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include + +namespace dflash::common { + +struct Qwen35LayerSplitAdapterConfig { + const char * target_path = nullptr; + const char * draft_path = nullptr; + DevicePlacement device; + int draft_gpu = 0; + RemoteDraftConfig remote_draft; + + int fa_window = 2048; + int kq_stride_pad = 32; + int draft_ctx_max = 4096; + int max_verify_tokens = DFLASH27B_DRAFT_BLOCK_SIZE; + bool run_dflash = false; +}; + +class Qwen35LayerSplitAdapter : public LayerSplitAdapter { +public: + explicit Qwen35LayerSplitAdapter(const Qwen35LayerSplitAdapterConfig & cfg); + ~Qwen35LayerSplitAdapter() override; + + Qwen35LayerSplitAdapter(const Qwen35LayerSplitAdapter &) = delete; + Qwen35LayerSplitAdapter & operator=(const Qwen35LayerSplitAdapter &) = delete; + + const char * name() const override { return "qwen35"; } + bool init() override; + int max_context() const override { return cfg_.device.max_ctx; } + + void begin_request(const GenerateRequest & req) override; + void reset_request_state() override; + bool prefill(const std::vector & prompt, + int base_pos, int & last_tok) override; + bool decode_ar(int last_tok, int committed, int n_gen, + std::vector & out_tokens, + const DaemonIO & io) override; + + bool can_dflash_decode() const override; + bool decode_dflash(const std::vector & prompt, int base_pos, + int last_tok, int n_gen, std::vector & out_tokens, + const DaemonIO & io) override; + + ModelBackend::CompressResult + compress(const ModelBackend::CompressRequest & req) override; + const char * default_compress_drafter_path() const override; + void free_drafter() override; + + bool snapshot_save(int slot) override; + void snapshot_free(int slot) override; + bool snapshot_used(int slot) const override; + int snapshot_cur_pos(int slot) const override; + bool snapshot_restore(int slot) override; + int current_last_token() const override; + + bool supports_dflash_spec_decode() const override { return true; } + DFlashTarget * dflash_target() override; + bool supports_remote_draft() const override { return true; } + + void shutdown() override; + +private: + bool load_draft(); + bool snapshot_slot_valid(int slot) const; + + Qwen35LayerSplitAdapterConfig cfg_; + std::vector shards_; + ggml_backend_t draft_backend_ = nullptr; + bool draft_backend_owned_ = false; + DraftWeights draft_weights_; + DraftFeatureMirror feature_ring_; + DFlashDraftIpcClient remote_draft_; + StepGraph draft_sg_; + StepGraph proj_sg_; + DrafterContext pflash_drafter_; + bool pflash_drafter_loaded_ = false; + static constexpr int PREFIX_SLOTS = ModelBackend::kMaxSlots; + std::vector> prefix_snapshots_; + std::vector snapshot_backends_; + + SamplerCfg sampler_; + std::mt19937_64 sampler_rng_{std::random_device{}()}; + std::unique_ptr dflash_target_; +}; + +} // namespace dflash::common diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index a010d1674..a4ccc5a59 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -1304,6 +1304,7 @@ bool snapshot_target_cache(const TargetWeights & w, for (int i = 0; i < n_full_attn; i++) { ggml_tensor * sk = cache.attn_k[i]; ggml_tensor * sv = cache.attn_v[i]; + if (!sk || !sv) continue; ggml_tensor * K = ggml_new_tensor_3d(snap.ctx, sk->type, sk->ne[0], snap_pos, sk->ne[2]); ggml_tensor * V = ggml_new_tensor_3d(snap.ctx, sv->type, sv->ne[0], snap_pos, sv->ne[2]); char name[64]; @@ -1317,6 +1318,7 @@ bool snapshot_target_cache(const TargetWeights & w, for (int i = 0; i < n_delta; i++) { ggml_tensor * ss = cache.ssm_state[i]; ggml_tensor * cs = cache.conv_state[i]; + if (!ss || !cs) continue; ggml_tensor * S = ggml_new_tensor_3d(snap.ctx, ss->type, ss->ne[0], ss->ne[1], ss->ne[2]); ggml_tensor * C = ggml_new_tensor_2d(snap.ctx, cs->type, cs->ne[0], cs->ne[1]); char name[64]; @@ -1327,11 +1329,13 @@ bool snapshot_target_cache(const TargetWeights & w, } // Right-sized target_feat: [fc_in, min(snap_pos, target_feat_cap)] - { + if (cache.target_feat) { ggml_tensor * tf = cache.target_feat; const int feat_len = std::min(snap_pos, cache.target_feat_cap); snap.target_feat_snap = ggml_new_tensor_2d(snap.ctx, tf->type, tf->ne[0], feat_len); ggml_set_name(snap.target_feat_snap, "snap_target_feat"); + } else { + snap.target_feat_snap = nullptr; } snap.buf = ggml_backend_alloc_ctx_tensors(snap.ctx, backend); @@ -1358,6 +1362,7 @@ bool snapshot_target_cache(const TargetWeights & w, ggml_tensor * dk = snap.attn_k_snap[i]; ggml_tensor * sv = cache.attn_v[i]; ggml_tensor * dv = snap.attn_v_snap[i]; + if (!sk || !dk || !sv || !dv) continue; const size_t k_strip = (size_t)snap_pos * sk->nb[1]; const size_t v_strip = (size_t)snap_pos * sv->nb[1]; for (int kh = 0; kh < (int)sk->ne[2]; kh++) { @@ -1374,12 +1379,16 @@ bool snapshot_target_cache(const TargetWeights & w, // SSM/conv: full copy (fixed-size, same shapes). for (int i = 0; i < n_delta; i++) { + if (!cache.ssm_state[i] || !snap.ssm_state_snap[i] || + !cache.conv_state[i] || !snap.conv_state_snap[i]) { + continue; + } ggml_backend_tensor_copy(cache.ssm_state[i], snap.ssm_state_snap[i]); ggml_backend_tensor_copy(cache.conv_state[i], snap.conv_state_snap[i]); } // target_feat: partial copy of first min(snap_pos, cap) rows. - { + if (cache.target_feat && snap.target_feat_snap) { const size_t feat_nbytes = ggml_nbytes(snap.target_feat_snap); ggml_backend_tensor_get(cache.target_feat, snap.target_feat_snap->data, 0, feat_nbytes); } @@ -1428,6 +1437,11 @@ bool restore_target_cache(const PrefixSnapshot & snap, TargetCache & cache) { ggml_tensor * dk = cache.attn_k[i]; ggml_tensor * sv = snap.attn_v_snap[i]; ggml_tensor * dv = cache.attn_v[i]; + if ((!sk || !sv) != (!dk || !dv)) { + set_last_error("restore_target_cache: KV shard layout mismatch"); + return false; + } + if (!sk || !dk || !sv || !dv) continue; const size_t k_strip = (size_t)snap_pos * sk->nb[1]; const size_t v_strip = (size_t)snap_pos * sv->nb[1]; for (int kh = 0; kh < (int)sk->ne[2]; kh++) { @@ -1444,12 +1458,21 @@ bool restore_target_cache(const PrefixSnapshot & snap, TargetCache & cache) { // SSM/conv: full copy (fixed-size). for (int i = 0; i < n_delta; i++) { + if ((!snap.ssm_state_snap[i] || !snap.conv_state_snap[i]) != + (!cache.ssm_state[i] || !cache.conv_state[i])) { + set_last_error("restore_target_cache: recurrent shard layout mismatch"); + return false; + } + if (!snap.ssm_state_snap[i] || !cache.ssm_state[i] || + !snap.conv_state_snap[i] || !cache.conv_state[i]) { + continue; + } ggml_backend_tensor_copy(snap.ssm_state_snap[i], cache.ssm_state[i]); ggml_backend_tensor_copy(snap.conv_state_snap[i], cache.conv_state[i]); } // target_feat: partial copy of stored rows. - { + if (cache.target_feat && snap.target_feat_snap) { const size_t feat_nbytes = ggml_nbytes(snap.target_feat_snap); ggml_backend_tensor_set(cache.target_feat, snap.target_feat_snap->data, 0, feat_nbytes); } diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index c34f94c94..fc06e2c1b 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -16,6 +16,7 @@ #include "model_card.h" #include "common/backend_factory.h" #include "common/gguf_inspect.h" +#include "common/layer_split_utils.h" #include "common/peer_access.h" #include "placement/pflash_placement.h" @@ -114,16 +115,25 @@ static bool validate_server_placement(const BackendArgs & bargs, placement_backend_name(compiled)); return false; } - if (!bargs.device.layer_split_gpus.empty()) { + if (!bargs.device.is_layer_split() && !bargs.device.layer_split_weights.empty()) { std::fprintf(stderr, - "[server] target layer split is not implemented in the native " - "server yet (--target-devices was provided)\n"); + "[server] --target-layer-split requires --target-devices\n"); return false; } - if (!bargs.device.layer_split_weights.empty()) { + if (bargs.device.is_layer_split()) { + const std::string placement_error = + validate_device_placement(bargs.device, /*device_count=*/-1); + if (!placement_error.empty()) { + std::fprintf(stderr, "[server] bad target layer split: %s\n", + placement_error.c_str()); + return false; + } + } + if (bargs.device.is_layer_split() && target != compiled) { std::fprintf(stderr, - "[server] --target-layer-split requires native target layer split " - "support, which is not implemented yet\n"); + "[server] target layer split must use this binary's compiled " + "backend (target=%s compiled=%s)\n", + placement_backend_name(target), placement_backend_name(compiled)); return false; } return true; @@ -451,6 +461,9 @@ int main(int argc, char ** argv) { sconfig.pflash_drafter_gpu = pflash_placement.drafter_gpu; sconfig.pflash_remote_drafter = pflash_placement.remote_drafter; sconfig.pflash_remote = pflash_placement.remote; + if (bargs.device.is_layer_split()) { + sconfig.disk_cache_dir.clear(); + } // ── Apply environment defaults ───────────────────────────────────── // Explicit --cache-type-k/v override via env vars. diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 0cbe7bfae..c3b803a6a 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -19,6 +19,10 @@ #include "common/sampler.h" #include "common/backend_ipc.h" #include "placement/pflash_placement.h" +#include "common/io_utils.h" +#include "placement/placement_config.h" +#include "common/layer_split_backend.h" +#include "common/layer_split_utils.h" #include #include @@ -1184,6 +1188,196 @@ static void test_normalize_responses_tool_followup_messages() { } // ═══════════════════════════════════════════════════════════════════════ +// Placement config tests +// ═══════════════════════════════════════════════════════════════════════ + +static void test_parse_target_device_list_same_backend() { + DevicePlacement placement; + TEST_ASSERT(parse_placement_device_list("cuda:0,cuda:1", placement)); + TEST_ASSERT(placement.backend == PlacementBackend::Cuda); + TEST_ASSERT(placement.gpu == 0); + TEST_ASSERT(placement.is_layer_split()); + TEST_ASSERT(placement.layer_split_gpus.size() == 2); + TEST_ASSERT(placement.layer_split_gpus[0] == 0); + TEST_ASSERT(placement.layer_split_gpus[1] == 1); + TEST_ASSERT(placement.layer_split_weights.empty()); +} + +static void test_parse_target_device_list_rejects_mixed_backend() { + DevicePlacement placement; + TEST_ASSERT(!parse_placement_device_list("cuda:0,hip:1", placement)); +} + +static void test_parse_target_device_list_single_gpu_is_not_layer_split() { + DevicePlacement placement; + TEST_ASSERT(parse_placement_device_list("hip:2", placement)); + TEST_ASSERT(placement.backend == PlacementBackend::Hip); + TEST_ASSERT(placement.gpu == 2); + TEST_ASSERT(!placement.is_layer_split()); + TEST_ASSERT(placement.layer_split_gpus.empty()); +} + +static void test_validate_layer_split_weights_shape() { + DevicePlacement placement; + TEST_ASSERT(parse_placement_device_list("cuda:0,cuda:1", placement)); + + placement.layer_split_weights = {1.0}; + TEST_ASSERT(!validate_device_placement(placement, -1).empty()); + + placement.layer_split_weights = {1.0, 0.0}; + TEST_ASSERT(!validate_device_placement(placement, -1).empty()); + + placement.layer_split_weights = {1.0, 2.0}; + TEST_ASSERT(validate_device_placement(placement, -1).empty()); +} + +struct MockLayerSplitAdapter : LayerSplitAdapter { + int max_ctx = 128; + bool reset_called = false; + int saved_slot = -1; + int saved_pos = 0; + int restored_slot = -1; + int current_pos = 0; + int current_last = -1; + std::vector prefill_bases; + std::vector prefill_sizes; + std::vector emitted_tokens; + bool dflash_enabled = false; + bool dflash_called = false; + ModelBackend::CompressRequest last_compress_req; + + const char * name() const override { return "mock"; } + bool init() override { return true; } + int max_context() const override { return max_ctx; } + void reset_request_state() override { + reset_called = true; + current_pos = 0; + current_last = -1; + } + bool prefill(const std::vector & prompt, + int base_pos, int & last_tok) override { + prefill_bases.push_back(base_pos); + prefill_sizes.push_back((int)prompt.size()); + current_pos = base_pos + (int)prompt.size(); + current_last = prompt.empty() ? current_last : prompt.back(); + last_tok = current_last; + return true; + } + bool decode_ar(int last_tok, int committed, int n_gen, + std::vector & out_tokens, + const DaemonIO & io) override { + TEST_ASSERT(committed == current_pos); + for (int i = 0; i < n_gen; ++i) { + int32_t tok = last_tok + i + 1; + out_tokens.push_back(tok); + emitted_tokens.push_back(tok); + io.emit(tok); + } + io.emit(-1); + return true; + } + bool can_dflash_decode() const override { return dflash_enabled; } + bool decode_dflash(const std::vector &, int, int, int, + std::vector &, const DaemonIO &) override { + dflash_called = true; + return false; + } + void free_drafter() override {} + bool snapshot_save(int slot) override { + saved_slot = slot; + saved_pos = current_pos; + return true; + } + bool snapshot_used(int slot) const override { + return slot == saved_slot && saved_pos > 0; + } + int snapshot_cur_pos(int slot) const override { + return snapshot_used(slot) ? saved_pos : 0; + } + bool snapshot_restore(int slot) override { + if (!snapshot_used(slot)) return false; + restored_slot = slot; + current_pos = saved_pos; + current_last = saved_pos; + return true; + } + int current_last_token() const override { return current_last; } + const char * default_compress_drafter_path() const override { + return "/tmp/default-layer-split-drafter.gguf"; + } + ModelBackend::CompressResult + compress(const ModelBackend::CompressRequest & req) override { + last_compress_req = req; + ModelBackend::CompressResult result; + result.ok = true; + result.compressed_ids = {77, 88}; + return result; + } + void shutdown() override {} +}; + +static void test_layer_split_backend_inline_snapshot_and_restore_delta() { + auto * raw = new MockLayerSplitAdapter(); + LayerSplitBackend backend{std::unique_ptr(raw)}; + + GenerateRequest req; + req.prompt = {10, 11, 12, 13}; + req.n_gen = 1; + req.snap_slot = 2; + req.snap_pos = 3; + DaemonIO io; + GenerateResult result = backend.generate(req, io); + + TEST_ASSERT(result.ok); + TEST_ASSERT(raw->reset_called); + TEST_ASSERT(raw->saved_slot == 2); + TEST_ASSERT(raw->saved_pos == 3); + TEST_ASSERT(raw->prefill_bases.size() == 2); + TEST_ASSERT(raw->prefill_bases[0] == 0); + TEST_ASSERT(raw->prefill_sizes[0] == 3); + TEST_ASSERT(raw->prefill_bases[1] == 3); + TEST_ASSERT(raw->prefill_sizes[1] == 1); + TEST_ASSERT(backend.snapshot_used(2)); + TEST_ASSERT(backend.snapshot_cur_pos(2) == 3); + + raw->reset_called = false; + raw->prefill_bases.clear(); + raw->prefill_sizes.clear(); + raw->dflash_enabled = true; + GenerateRequest restore_req; + restore_req.prompt = {10, 11, 12, 99}; + restore_req.n_gen = 1; + GenerateResult restored = backend.restore_and_generate(2, restore_req, io); + + TEST_ASSERT(restored.ok); + TEST_ASSERT(!raw->dflash_called); + TEST_ASSERT(raw->restored_slot == 2); + TEST_ASSERT(!raw->reset_called); + TEST_ASSERT(raw->prefill_bases.size() == 1); + TEST_ASSERT(raw->prefill_bases[0] == 3); + TEST_ASSERT(raw->prefill_sizes[0] == 1); +} + +static void test_layer_split_compress_nopark_uses_default_drafter_path() { + const std::string ids_path = "/tmp/dflash_test_layer_split_compress_ids.bin"; + unlink(ids_path.c_str()); + TEST_ASSERT(write_int32_file(ids_path, {1, 2, 3, 4})); + + auto * raw = new MockLayerSplitAdapter(); + LayerSplitBackend backend{std::unique_ptr(raw)}; + DaemonIO io; + + const std::string cmd = "compress " + ids_path + " 250 nopark"; + TEST_ASSERT(backend.handle_compress(cmd, io)); + TEST_ASSERT(raw->last_compress_req.skip_park); + TEST_ASSERT(std::abs(raw->last_compress_req.keep_ratio - 0.25f) < 1e-5f); + TEST_ASSERT(raw->last_compress_req.input_ids.size() == 4); + TEST_ASSERT(raw->last_compress_req.drafter_path == + "/tmp/default-layer-split-drafter.gguf"); + + unlink(ids_path.c_str()); +} + // Disk Prefix Cache Tests // ═══════════════════════════════════════════════════════════════════════ @@ -2308,6 +2502,14 @@ int main() { RUN_TEST(test_jinja_render_bad_tools_json_throws); RUN_TEST(test_normalize_responses_tool_followup_messages); + std::fprintf(stderr, "\n── Placement config ──\n"); + RUN_TEST(test_parse_target_device_list_same_backend); + RUN_TEST(test_parse_target_device_list_rejects_mixed_backend); + RUN_TEST(test_parse_target_device_list_single_gpu_is_not_layer_split); + RUN_TEST(test_validate_layer_split_weights_shape); + RUN_TEST(test_layer_split_backend_inline_snapshot_and_restore_delta); + RUN_TEST(test_layer_split_compress_nopark_uses_default_drafter_path); + std::fprintf(stderr, "\n── Disk prefix cache ──\n"); RUN_TEST(test_disk_cache_config_defaults); RUN_TEST(test_disk_cache_disabled_when_no_dir); From 599217dbd97bf49db10bf0d9b2b23a1a10738a89 Mon Sep 17 00:00:00 2001 From: weicj Date: Sun, 24 May 2026 03:59:55 +0800 Subject: [PATCH 2/7] refactor(server): generalize target layer-split adapter path --- .../src/common/dflash_layer_split_runtime.h | 12 +-- server/src/common/layer_split_utils.cpp | 75 +++++++++++++++- server/src/common/layer_split_utils.h | 67 ++++++++++++++- server/src/internal.h | 1 + server/src/qwen35/gguf_target_loader.cpp | 1 + server/src/qwen35/layer_split_daemon.cpp | 8 +- server/src/qwen35/layer_split_daemon.h | 6 +- server/src/qwen35/layer_split_daemon_loop.cpp | 50 ++++------- server/src/qwen35/layer_split_daemon_loop.h | 2 +- server/src/qwen35/layer_split_forward.cpp | 12 +-- server/src/qwen35/layer_split_forward.h | 6 +- server/src/qwen35/layer_split_types.h | 26 ++---- server/src/qwen35/qwen35_layer_split.h | 2 +- .../src/qwen35/qwen35_layer_split_adapter.cpp | 57 ++++-------- .../src/qwen35/qwen35_layer_split_adapter.h | 2 +- .../qwen35_layer_split_dflash_target.cpp | 4 +- .../qwen35/qwen35_layer_split_dflash_target.h | 8 +- server/test/test_dflash.cpp | 86 ++++++++----------- 18 files changed, 245 insertions(+), 180 deletions(-) diff --git a/server/src/common/dflash_layer_split_runtime.h b/server/src/common/dflash_layer_split_runtime.h index 5c539989f..af05d5ad6 100644 --- a/server/src/common/dflash_layer_split_runtime.h +++ b/server/src/common/dflash_layer_split_runtime.h @@ -1,11 +1,11 @@ // dflash_layer_split_runtime.h — target-agnostic runtime types for the -// DFlash layer-split pipeline. +// target layer-split pipeline. // -// Hosts the small pieces that are reused by every architecture's layer-split -// driver: a runtime-configuration struct (replaces former globals) and an -// activation double-buffer used to ferry hidden states between shards. -// Architecture-specific shard layouts (e.g. qwen35's TargetLayerSplitShard -// that embeds TargetWeights/TargetCache) live in their own headers. +// Hosts the small runtime pieces reused by layer-split drivers: a +// runtime-configuration struct and the activation double-buffer used to ferry +// hidden states between shards. Shared placement/load-plan/shard metadata lives +// in common/layer_split_utils.h; architecture-specific shard payloads keep their +// own weights/cache/graph types. #pragma once diff --git a/server/src/common/layer_split_utils.cpp b/server/src/common/layer_split_utils.cpp index f69501bf0..e1f9f3f42 100644 --- a/server/src/common/layer_split_utils.cpp +++ b/server/src/common/layer_split_utils.cpp @@ -1,17 +1,22 @@ #include "layer_split_utils.h" +#include "common/peer_access.h" +#include "common/snapshot_backend.h" +#include "ggml-cuda.h" + #include #include +#include #include namespace dflash::common { -std::vector> compute_layer_ranges( +std::vector compute_layer_ranges( int n_layer, int n_gpus, const std::vector & weights) { - std::vector> ranges; + std::vector ranges; if (n_layer <= 0 || n_gpus <= 0 || n_gpus > n_layer) return ranges; std::vector w = weights; @@ -39,6 +44,72 @@ std::vector> compute_layer_ranges( return ranges; } +bool init_layer_split_shard_metas( + std::vector shards, + const std::vector & gpus, + const std::vector & ranges, + const char * log_prefix) { + if (shards.size() != gpus.size() || shards.size() != ranges.size()) return false; + const char * prefix = log_prefix ? log_prefix : "target-split"; + for (size_t i = 0; i < shards.size(); ++i) { + auto * shard = shards[i]; + if (!shard) return false; + shard->gpu = gpus[i]; + shard->layer_begin = ranges[i].begin; + shard->layer_end = ranges[i].end; + shard->backend = ggml_backend_cuda_init(shard->gpu); + if (!shard->backend) { + std::fprintf(stderr, "[%s] backend init failed gpu=%d\n", + prefix, shard->gpu); + return false; + } + } + return true; +} + +bool enable_layer_split_peer_access( + const std::vector & gpus, + bool peer_access) { + if (!peer_access) return true; + for (size_t i = 0; i < gpus.size(); ++i) { + for (size_t j = i + 1; j < gpus.size(); ++j) { + (void)enable_peer_access_pair(gpus[i], gpus[j]); + } + } + return true; +} + +bool init_layer_split_snapshot_backends( + const std::vector & shards, + std::vector & snapshot_backends, + const char * log_prefix) { + const char * prefix = log_prefix ? log_prefix : "target-split"; + snapshot_backends.assign(shards.size(), nullptr); + for (size_t i = 0; i < shards.size(); ++i) { + const auto * shard = shards[i]; + if (!shard || !shard->backend) return false; + snapshot_backends[i] = create_snapshot_backend(shard->backend); + if (!snapshot_backends[i]) { + std::fprintf(stderr, + "[%s] snapshot backend init failed gpu=%d\n", + prefix, shard->gpu); + return false; + } + } + return true; +} + +void free_layer_split_snapshot_backends( + const std::vector & shards, + std::vector & snapshot_backends) { + const size_t n = std::min(shards.size(), snapshot_backends.size()); + for (size_t i = 0; i < n; ++i) { + if (!shards[i]) continue; + free_snapshot_backend(snapshot_backends[i], shards[i]->backend); + } + snapshot_backends.clear(); +} + std::string validate_device_placement( const DevicePlacement & dp, int device_count) diff --git a/server/src/common/layer_split_utils.h b/server/src/common/layer_split_utils.h index 98ff422af..447c41ccc 100644 --- a/server/src/common/layer_split_utils.h +++ b/server/src/common/layer_split_utils.h @@ -8,20 +8,83 @@ #include "placement/placement_config.h" #include -#include #include +#include "ggml-backend.h" + namespace dflash::common { +struct LayerSplitRange { + int begin = 0; + int end = 0; +}; + +struct LayerSplitShardMeta { + int gpu = 0; + int layer_begin = 0; + int layer_end = 0; + ggml_backend_t backend = nullptr; +}; + +inline TargetLoadPlan make_layer_split_load_plan( + const LayerSplitShardMeta & shard, + bool is_last_shard) { + TargetLoadPlan plan; + plan.layer_begin = shard.layer_begin; + plan.layer_end = shard.layer_end; + plan.load_output = is_last_shard; + return plan; +} + +template +std::vector layer_split_shard_metas( + std::vector & shards) { + std::vector metas; + metas.reserve(shards.size()); + for (auto & shard : shards) { + metas.push_back(&shard); + } + return metas; +} + +template +Shard * find_layer_split_shard(std::vector & shards, int layer_idx) { + for (auto & shard : shards) { + if (layer_idx >= shard.layer_begin && layer_idx < shard.layer_end) { + return &shard; + } + } + return nullptr; +} + // Compute [begin, end) layer ranges for each GPU shard. // If weights is empty, splits layers equally. // If weights has entries, distributes proportionally (at least 1 layer per GPU). // Returns empty vector on error (n_layer <= 0 or n_gpus <= 0). -std::vector> compute_layer_ranges( +std::vector compute_layer_ranges( int n_layer, int n_gpus, const std::vector & weights); +bool init_layer_split_shard_metas( + std::vector shards, + const std::vector & gpus, + const std::vector & ranges, + const char * log_prefix); + +bool enable_layer_split_peer_access( + const std::vector & gpus, + bool peer_access); + +bool init_layer_split_snapshot_backends( + const std::vector & shards, + std::vector & snapshot_backends, + const char * log_prefix); + +void free_layer_split_snapshot_backends( + const std::vector & shards, + std::vector & snapshot_backends); + // Validate a DevicePlacement against system constraints. // If device_count is negative, only validates structural constraints that do // not require querying the runtime-visible GPU count. diff --git a/server/src/internal.h b/server/src/internal.h index f9a890ff2..adeab6321 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -23,6 +23,7 @@ #include "ggml-backend.h" #include "gguf.h" +#include "common/layer_split_utils.h" #include "dflash27b.h" namespace dflash::common { diff --git a/server/src/qwen35/gguf_target_loader.cpp b/server/src/qwen35/gguf_target_loader.cpp index bafa57bad..116ddafc0 100644 --- a/server/src/qwen35/gguf_target_loader.cpp +++ b/server/src/qwen35/gguf_target_loader.cpp @@ -44,6 +44,7 @@ // tensor's bytes from the mmap'd file. #include "internal.h" +#include "common/layer_split_utils.h" #include #include diff --git a/server/src/qwen35/layer_split_daemon.cpp b/server/src/qwen35/layer_split_daemon.cpp index c8c107bf9..d0a833174 100644 --- a/server/src/qwen35/layer_split_daemon.cpp +++ b/server/src/qwen35/layer_split_daemon.cpp @@ -14,8 +14,8 @@ namespace dflash::common { -bool run_target_layer_split_request( - std::vector & shards, +bool run_qwen35_layer_split_request( + std::vector & shards, DraftWeights * draft_weights, ggml_backend_t draft_backend, int draft_gpu, @@ -41,7 +41,7 @@ bool run_target_layer_split_request( ubatch = std::max(1, std::atoi(s)); } int last_tok = -1; - if (!run_target_layer_split_forward(shards, shards.front().weights, + if (!run_qwen35_layer_split_forward(shards, shards.front().weights, prompt, 0, ubatch, last_tok, kq_stride_pad, fa_window, feature_ring)) { @@ -66,7 +66,7 @@ bool run_target_layer_split_request( for (; generated < n_gen; generated++) { std::vector one(1, last_tok); int next_tok = -1; - if (!run_target_layer_split_forward(shards, shards.front().weights, + if (!run_qwen35_layer_split_forward(shards, shards.front().weights, one, (int)out_all.size(), 1, next_tok, kq_stride_pad, fa_window, feature_ring)) { diff --git a/server/src/qwen35/layer_split_daemon.h b/server/src/qwen35/layer_split_daemon.h index 3d87a3313..13777f18b 100644 --- a/server/src/qwen35/layer_split_daemon.h +++ b/server/src/qwen35/layer_split_daemon.h @@ -1,6 +1,6 @@ // layer_split_daemon.h — Layer-split request handler for qwen35 daemon mode. // -// run_target_layer_split_request() handles a single inference request: +// run_qwen35_layer_split_request() handles a single inference request: // prefill → (optional spec-decode or AR decode) → output. #pragma once @@ -22,8 +22,8 @@ namespace dflash::common { // Runs prefill, then either spec-decode (if run_dflash && draft available) // or plain AR decode. Emits tokens to stream_fd and optionally writes // the full sequence to out_path. -bool run_target_layer_split_request( - std::vector & shards, +bool run_qwen35_layer_split_request( + std::vector & shards, DraftWeights * draft_weights, ggml_backend_t draft_backend, int draft_gpu, diff --git a/server/src/qwen35/layer_split_daemon_loop.cpp b/server/src/qwen35/layer_split_daemon_loop.cpp index 6790f5510..be2d4d535 100644 --- a/server/src/qwen35/layer_split_daemon_loop.cpp +++ b/server/src/qwen35/layer_split_daemon_loop.cpp @@ -2,10 +2,9 @@ #include "layer_split_daemon_loop.h" #include "layer_split_types.h" -#include "layer_split_daemon.h" // run_target_layer_split_request -#include "layer_split_forward.h" // free_target_layer_split_shards +#include "layer_split_daemon.h" // run_qwen35_layer_split_request +#include "layer_split_forward.h" // free_qwen35_layer_split_shards #include "dflash_feature_ring.h" -#include "peer_access.h" #include "common/io_utils.h" #include "common/sampler.h" #include "common/layer_split_utils.h" @@ -38,34 +37,19 @@ int run_layer_split_daemon(const LayerSplitDaemonConfig & cfg) { } // Initialize shards. - std::vector shards(cfg.target_gpus.size()); - for (size_t i = 0; i < cfg.target_gpus.size(); i++) { - shards[i].gpu = cfg.target_gpus[i]; - shards[i].layer_begin = ranges[i].first; - shards[i].layer_end = ranges[i].second; - shards[i].backend = ggml_backend_cuda_init(shards[i].gpu); - if (!shards[i].backend) { - std::fprintf(stderr, "target-split cuda init failed for gpu %d\n", shards[i].gpu); - free_target_layer_split_shards(shards); - return 1; - } - } - - // Enable peer access between all GPU pairs. - for (size_t i = 0; i < cfg.target_gpus.size(); i++) { - for (size_t j = i + 1; j < cfg.target_gpus.size(); j++) { - if (cfg.peer_access) { - (void)enable_peer_access_pair(cfg.target_gpus[i], cfg.target_gpus[j]); - } - } + std::vector shards(cfg.target_gpus.size()); + auto shard_metas = layer_split_shard_metas(shards); + if (!init_layer_split_shard_metas( + shard_metas, cfg.target_gpus, ranges, "target-split")) { + free_qwen35_layer_split_shards(shards); + return 1; } + (void)enable_layer_split_peer_access(cfg.target_gpus, cfg.peer_access); // Load partial target weights + caches. for (auto & shard : shards) { - TargetLoadPlan plan; - plan.layer_begin = shard.layer_begin; - plan.layer_end = shard.layer_end; - plan.load_output = (&shard == &shards.back()); + const TargetLoadPlan plan = + make_layer_split_load_plan(shard, &shard == &shards.back()); if (!load_target_gguf_partial(cfg.target_path, shard.backend, plan, shard.weights) || !create_target_cache_partial(shard.weights, cfg.max_ctx, cfg.max_verify_tokens, shard.backend, shard.cache, @@ -74,7 +58,7 @@ int run_layer_split_daemon(const LayerSplitDaemonConfig & cfg) { /*allocate_target_feat=*/false)) { std::fprintf(stderr, "target-split load/cache gpu=%d: %s\n", shard.gpu, dflash27b_last_error()); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } } @@ -91,7 +75,7 @@ int run_layer_split_daemon(const LayerSplitDaemonConfig & cfg) { if (!draft_backend) { draft_backend = ggml_backend_cuda_init(cfg.draft_gpu); if (!draft_backend) { - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } draft_backend_owned = true; @@ -105,7 +89,7 @@ int run_layer_split_daemon(const LayerSplitDaemonConfig & cfg) { std::fprintf(stderr, "target-split draft load gpu=%d: %s\n", cfg.draft_gpu, dflash27b_last_error()); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } const int cap = std::min(cfg.max_ctx, 4096); @@ -117,7 +101,7 @@ int run_layer_split_daemon(const LayerSplitDaemonConfig & cfg) { cfg.draft_gpu); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } } @@ -172,7 +156,7 @@ int run_layer_split_daemon(const LayerSplitDaemonConfig & cfg) { for (auto & shard : shards) { reset_target_cache(shard.cache); } - const bool ok = run_target_layer_split_request( + const bool ok = run_qwen35_layer_split_request( shards, cfg.load_draft ? &draft_weights : nullptr, draft_backend, @@ -188,7 +172,7 @@ int run_layer_split_daemon(const LayerSplitDaemonConfig & cfg) { draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned && draft_backend) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 0; } diff --git a/server/src/qwen35/layer_split_daemon_loop.h b/server/src/qwen35/layer_split_daemon_loop.h index 2215f00ba..75112280c 100644 --- a/server/src/qwen35/layer_split_daemon_loop.h +++ b/server/src/qwen35/layer_split_daemon_loop.h @@ -5,7 +5,7 @@ // owns a shard of layers and the activation tensor is forwarded between shards. // // Kept separate from daemon_loop.{h,cpp} because: -// - It manages multiple TargetLayerSplitShards instead of a single ModelBackend +// - It manages multiple Qwen35LayerSplitShards instead of a single ModelBackend // - Peer access / feature mirroring is shard-specific // - SNAPSHOT/RESTORE are not supported in layer-split mode diff --git a/server/src/qwen35/layer_split_forward.cpp b/server/src/qwen35/layer_split_forward.cpp index 1ee664301..d1ab66587 100644 --- a/server/src/qwen35/layer_split_forward.cpp +++ b/server/src/qwen35/layer_split_forward.cpp @@ -61,8 +61,8 @@ bool compute_target_split_argmax( return true; } -bool run_target_layer_split_forward( - std::vector & shards, +bool run_qwen35_layer_split_forward( + std::vector & shards, const TargetWeights & embed_source, const std::vector & tokens, int base_pos, @@ -104,11 +104,11 @@ bool run_target_layer_split_forward( } } - TargetLayerSplitShard * current_shard = &shards.front(); + Qwen35LayerSplitShard * current_shard = &shards.front(); std::vector mask_buf; std::vector pos_buf; for (int il = 0; il < embed_source.n_layer; il++) { - TargetLayerSplitShard * shard = find_target_shard(shards, il); + Qwen35LayerSplitShard * shard = find_layer_split_shard(shards, il); if (!shard) { std::fprintf(stderr, "target-split missing owner for layer %d\n", il); activation_pair_free(acts); @@ -204,7 +204,7 @@ bool run_target_layer_split_forward( StepGraph final_sg; std::vector argmax_tokens; - TargetLayerSplitShard & last_shard = shards.back(); + Qwen35LayerSplitShard & last_shard = shards.back(); const bool need_all_argmax = argmax_out != nullptr; const int argmax_offset = need_all_argmax ? 0 : (n_tokens_total - 1); const int argmax_count = need_all_argmax ? n_tokens_total : 1; @@ -224,7 +224,7 @@ bool run_target_layer_split_forward( return true; } -void free_target_layer_split_shards(std::vector & shards) { +void free_qwen35_layer_split_shards(std::vector & shards) { for (auto & shard : shards) { step_graph_destroy(shard.layer_graph); free_target_cache(shard.cache); diff --git a/server/src/qwen35/layer_split_forward.h b/server/src/qwen35/layer_split_forward.h index 7bf0b9222..c04680fe4 100644 --- a/server/src/qwen35/layer_split_forward.h +++ b/server/src/qwen35/layer_split_forward.h @@ -35,8 +35,8 @@ bool compute_target_split_argmax( // Run a full forward pass through all shards, writing K/V into each shard's // cache. Returns the argmax of the last token in `last_tok`. // Optionally captures features into `feature_ring` / remote draft. -bool run_target_layer_split_forward( - std::vector & shards, +bool run_qwen35_layer_split_forward( + std::vector & shards, const TargetWeights & embed_source, const std::vector & tokens, int base_pos, @@ -50,6 +50,6 @@ bool run_target_layer_split_forward( DFlashDraftIpcClient * remote_draft = nullptr); // Free all shards (weights, cache, backend). -void free_target_layer_split_shards(std::vector & shards); +void free_qwen35_layer_split_shards(std::vector & shards); } // namespace dflash::common diff --git a/server/src/qwen35/layer_split_types.h b/server/src/qwen35/layer_split_types.h index 3920d6f3c..6e222d767 100644 --- a/server/src/qwen35/layer_split_types.h +++ b/server/src/qwen35/layer_split_types.h @@ -1,13 +1,14 @@ // layer_split_types.h — qwen35 layer-split shard types. // -// Generic split-runtime helpers (LayerSplitRuntimeConfig, ActivationPair) -// live in `common/dflash_layer_split_runtime.h`. This header keeps the -// qwen35-specific shard layout that embeds TargetWeights/TargetCache. +// Shared layer-split metadata lives in common/layer_split_utils.h. This header +// keeps only the qwen35-specific shard payload: TargetWeights, TargetCache, and +// the qwen35 step graph. #pragma once #include "internal.h" #include "step_graph.h" +#include "common/layer_split_utils.h" #include "dflash_layer_split_runtime.h" #include "ggml.h" @@ -15,30 +16,15 @@ #include "ggml-backend.h" #include -#include namespace dflash::common { -// ── Per-GPU shard for layer-split target ──────────────────────────── +// ── Per-GPU qwen35 shard for layer-split target ───────────────────── -struct TargetLayerSplitShard { - int gpu = 0; - int layer_begin = 0; - int layer_end = 0; - ggml_backend_t backend = nullptr; +struct Qwen35LayerSplitShard : LayerSplitShardMeta { TargetWeights weights; TargetCache cache; StepGraph layer_graph; }; -inline TargetLayerSplitShard * find_target_shard( - std::vector & shards, - int layer_idx) { - for (auto & shard : shards) { - if (layer_idx >= shard.layer_begin && layer_idx < shard.layer_end) - return &shard; - } - return nullptr; -} - } // namespace dflash::common diff --git a/server/src/qwen35/qwen35_layer_split.h b/server/src/qwen35/qwen35_layer_split.h index 2893fb82b..9e789e40e 100644 --- a/server/src/qwen35/qwen35_layer_split.h +++ b/server/src/qwen35/qwen35_layer_split.h @@ -8,7 +8,7 @@ // (run_target_layer_split_daemon) because it depends on many helpers and // globals defined there. This header defines the DevicePlacement-based // args struct as the migration target. The implementation will move here -// once the helper functions (run_target_layer_split_forward, etc.) are +// once the helper functions (run_qwen35_layer_split_forward, etc.) are // extracted to src/qwen35/. #pragma once diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index 0d6dcc3d0..87e6d067c 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -5,9 +5,7 @@ #include "common/dflash_spec_decode.h" #include "common/gguf_inspect.h" #include "common/layer_split_utils.h" -#include "common/peer_access.h" #include "common/sampler.h" -#include "common/snapshot_backend.h" #include "qwen35/layer_split_forward.h" #include "qwen35/qwen35_layer_split_dflash_target.h" #include "qwen3/qwen3_drafter.h" @@ -50,44 +48,21 @@ bool Qwen35LayerSplitAdapter::init() { } shards_.resize(cfg_.device.layer_split_gpus.size()); - for (size_t i = 0; i < shards_.size(); ++i) { - auto & shard = shards_[i]; - shard.gpu = cfg_.device.layer_split_gpus[i]; - shard.layer_begin = ranges[i].first; - shard.layer_end = ranges[i].second; - shard.backend = ggml_backend_cuda_init(shard.gpu); - if (!shard.backend) { - std::fprintf(stderr, - "[target-split] backend init failed gpu=%d\n", shard.gpu); - return false; - } + auto shard_metas = layer_split_shard_metas(shards_); + if (!init_layer_split_shard_metas( + shard_metas, cfg_.device.layer_split_gpus, ranges, "target-split")) { + return false; } - if (cfg_.device.peer_access) { - for (size_t i = 0; i < cfg_.device.layer_split_gpus.size(); ++i) { - for (size_t j = i + 1; j < cfg_.device.layer_split_gpus.size(); ++j) { - (void)enable_peer_access_pair(cfg_.device.layer_split_gpus[i], - cfg_.device.layer_split_gpus[j]); - } - } - } + (void)enable_layer_split_peer_access( + cfg_.device.layer_split_gpus, cfg_.device.peer_access); - snapshot_backends_.resize(shards_.size(), nullptr); - for (size_t i = 0; i < shards_.size(); ++i) { - snapshot_backends_[i] = create_snapshot_backend(shards_[i].backend); - if (!snapshot_backends_[i]) { - std::fprintf(stderr, - "[target-split] snapshot backend init failed gpu=%d\n", - shards_[i].gpu); - return false; - } - } + if (!init_layer_split_snapshot_backends( + shard_metas, snapshot_backends_, "target-split")) return false; for (auto & shard : shards_) { - TargetLoadPlan plan; - plan.layer_begin = shard.layer_begin; - plan.layer_end = shard.layer_end; - plan.load_output = (&shard == &shards_.back()); + const TargetLoadPlan plan = + make_layer_split_load_plan(shard, &shard == &shards_.back()); if (!load_target_gguf_partial(cfg_.target_path, shard.backend, plan, shard.weights) || !create_target_cache_partial(shard.weights, cfg_.device.max_ctx, @@ -204,7 +179,7 @@ bool Qwen35LayerSplitAdapter::prefill(const std::vector & prompt, if (const char * s = std::getenv("DFLASH27B_PREFILL_UBATCH")) { ubatch = std::max(1, std::atoi(s)); } - return run_target_layer_split_forward( + return run_qwen35_layer_split_forward( shards_, shards_.front().weights, prompt, base_pos, ubatch, last_tok, cfg_.kq_stride_pad, /*fa_window=*/0, (cfg_.run_dflash && !remote_draft_.active()) ? &feature_ring_ : nullptr, @@ -297,7 +272,7 @@ bool Qwen35LayerSplitAdapter::decode_ar( for (int i = 1; i < n_gen; ++i) { std::vector one(1, last_tok); int next_tok = -1; - if (!run_target_layer_split_forward( + if (!run_qwen35_layer_split_forward( shards_, shards_.front().weights, one, committed, 1, next_tok, cfg_.kq_stride_pad, cfg_.fa_window, cfg_.run_dflash ? &feature_ring_ : nullptr)) { @@ -403,16 +378,14 @@ void Qwen35LayerSplitAdapter::shutdown() { for (auto & snap : slot) free_prefix_snapshot(snap); } prefix_snapshots_.clear(); - for (size_t i = 0; i < snapshot_backends_.size() && i < shards_.size(); ++i) { - free_snapshot_backend(snapshot_backends_[i], shards_[i].backend); - } - snapshot_backends_.clear(); + auto shard_metas = layer_split_shard_metas(shards_); + free_layer_split_snapshot_backends(shard_metas, snapshot_backends_); if (draft_backend_owned_ && draft_backend_) { ggml_backend_free(draft_backend_); } draft_backend_ = nullptr; draft_backend_owned_ = false; - free_target_layer_split_shards(shards_); + free_qwen35_layer_split_shards(shards_); } } // namespace dflash::common diff --git a/server/src/qwen35/qwen35_layer_split_adapter.h b/server/src/qwen35/qwen35_layer_split_adapter.h index f204c06ae..1feaf82dd 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.h +++ b/server/src/qwen35/qwen35_layer_split_adapter.h @@ -83,7 +83,7 @@ class Qwen35LayerSplitAdapter : public LayerSplitAdapter { bool snapshot_slot_valid(int slot) const; Qwen35LayerSplitAdapterConfig cfg_; - std::vector shards_; + std::vector shards_; ggml_backend_t draft_backend_ = nullptr; bool draft_backend_owned_ = false; DraftWeights draft_weights_; diff --git a/server/src/qwen35/qwen35_layer_split_dflash_target.cpp b/server/src/qwen35/qwen35_layer_split_dflash_target.cpp index 76daca64f..c602e7137 100644 --- a/server/src/qwen35/qwen35_layer_split_dflash_target.cpp +++ b/server/src/qwen35/qwen35_layer_split_dflash_target.cpp @@ -13,7 +13,7 @@ Qwen35LayerSplitDFlashTarget::~Qwen35LayerSplitDFlashTarget() { } Qwen35LayerSplitDFlashTarget::Qwen35LayerSplitDFlashTarget( - std::vector & shards, + std::vector & shards, DraftFeatureMirror * feature_ring, int kq_stride_pad, int fa_window, @@ -36,7 +36,7 @@ bool Qwen35LayerSplitDFlashTarget::verify_batch( int & last_tok, std::vector * all_argmax) { if (shards_.empty()) return false; - return run_target_layer_split_forward( + return run_qwen35_layer_split_forward( shards_, shards_.front().weights, tokens, base_pos, (int)tokens.size(), last_tok, kq_stride_pad_, fa_window_, feature_ring_, diff --git a/server/src/qwen35/qwen35_layer_split_dflash_target.h b/server/src/qwen35/qwen35_layer_split_dflash_target.h index f593009d8..40dc409d3 100644 --- a/server/src/qwen35/qwen35_layer_split_dflash_target.h +++ b/server/src/qwen35/qwen35_layer_split_dflash_target.h @@ -1,13 +1,13 @@ // Qwen35LayerSplitDFlashTarget — DFlashTarget adapter for qwen35 multi-GPU // layer-split inference. // -// Wraps a vector of TargetLayerSplitShard behind the generic DFlashTarget +// Wraps a vector of Qwen35LayerSplitShard behind the generic DFlashTarget // interface so the common spec-decode loop can drive layer-split verify // without knowing about shard layout, qwen attention parameters, or the // remote-draft IPC client. // // snapshot_kv / restore_kv iterate every shard. verify_batch delegates to -// run_target_layer_split_forward. project_hidden_to_tokens builds the +// run_qwen35_layer_split_forward. project_hidden_to_tokens builds the // LM-head graph on the back shard (where the lm_head weights live). #pragma once @@ -26,7 +26,7 @@ namespace dflash::common { class Qwen35LayerSplitDFlashTarget : public DFlashTarget { public: // All references/pointers are non-owning; caller controls lifetime. - Qwen35LayerSplitDFlashTarget(std::vector & shards, + Qwen35LayerSplitDFlashTarget(std::vector & shards, DraftFeatureMirror * feature_ring, int kq_stride_pad, int fa_window, @@ -56,7 +56,7 @@ class Qwen35LayerSplitDFlashTarget : public DFlashTarget { const std::vector & capture_layer_ids() const override; private: - std::vector & shards_; + std::vector & shards_; DraftFeatureMirror * feature_ring_; int kq_stride_pad_; int fa_window_; diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index 96db90834..26e51e4c2 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -169,8 +169,11 @@ using dflash::common::g_peer_access_opt_in; using dflash::common::g_peer_pair_ok_cache; using dflash::common::enable_peer_access_one_way; using dflash::common::enable_peer_access_pair; +using dflash::common::enable_layer_split_peer_access; using dflash::common::cross_device_peer_memcpy_ok; using dflash::common::copy_peer_async; +using dflash::common::init_layer_split_shard_metas; +using dflash::common::layer_split_shard_metas; using dflash::common::DraftFeatureMirror; using dflash::common::draft_feature_mirror_free; using dflash::common::draft_feature_mirror_init; @@ -190,11 +193,11 @@ using dflash::common::build_lm_head_projection_step; // ─── Layer split types — extracted to src/qwen35/layer_split_types.h ── #include "layer_split_types.h" using dflash::common::LayerSplitRuntimeConfig; -using dflash::common::TargetLayerSplitShard; +using dflash::common::Qwen35LayerSplitShard; using dflash::common::ActivationPair; using dflash::common::activation_pair_free; using dflash::common::activation_pair_init; -using dflash::common::find_target_shard; +using dflash::common::find_layer_split_shard; static bool parse_int_list(const char * text, std::vector & out) { out.clear(); @@ -251,8 +254,8 @@ using dflash::common::copy_feature_ring_range_to_tensor; // ─── Layer-split forward — extracted to src/qwen35/layer_split_forward.{h,cpp} ── #include "layer_split_forward.h" using dflash::common::compute_target_split_argmax; -using dflash::common::run_target_layer_split_forward; -using dflash::common::free_target_layer_split_shards; +using dflash::common::run_qwen35_layer_split_forward; +using dflash::common::free_qwen35_layer_split_shards; // ─── Speculative decode — generic loop in common/, qwen35 layer-split adapter. @@ -262,7 +265,7 @@ using dflash::common::is_eos_tok; // ─── Layer-split daemon — extracted to src/qwen35/layer_split_daemon.{h,cpp} ─ #include "layer_split_daemon.h" -using dflash::common::run_target_layer_split_request; +using dflash::common::run_qwen35_layer_split_request; static int run_target_layer_split_daemon( const char * target_path, @@ -330,39 +333,22 @@ static int run_target_layer_split_harness( target_gpus.size(), n_layer); return 2; } - std::vector shards; + std::vector shards; shards.resize(target_gpus.size()); - for (size_t i = 0; i < target_gpus.size(); i++) { - shards[i].gpu = target_gpus[i]; - shards[i].layer_begin = ranges[i].first; - shards[i].layer_end = ranges[i].second; - } - for (auto & shard : shards) { - shard.backend = ggml_backend_cuda_init(shard.gpu); - if (!shard.backend) { - std::fprintf(stderr, "target-split cuda init failed for gpu %d\n", shard.gpu); - free_target_layer_split_shards(shards); - return 1; - } - } - for (size_t i = 0; i < target_gpus.size(); i++) { - for (size_t j = i + 1; j < target_gpus.size(); j++) { - if (!enable_peer_access_pair(target_gpus[i], target_gpus[j])) { - std::fprintf(stderr, - "warning: CUDA peer access not fully enabled for target gpus %d,%d\n", - target_gpus[i], target_gpus[j]); - } - } + auto shard_metas = layer_split_shard_metas(shards); + if (!init_layer_split_shard_metas( + shard_metas, target_gpus, ranges, "target-split")) { + free_qwen35_layer_split_shards(shards); + return 1; } + (void)enable_layer_split_peer_access(target_gpus, peer_access); for (auto & shard : shards) { - TargetLoadPlan plan; - plan.layer_begin = shard.layer_begin; - plan.layer_end = shard.layer_end; - plan.load_output = (&shard == &shards.back()); + const TargetLoadPlan plan = + make_layer_split_load_plan(shard, &shard == &shards.back()); if (!load_target_gguf_partial(target_path, shard.backend, plan, shard.weights)) { std::fprintf(stderr, "target-split load gpu=%d: %s\n", shard.gpu, dflash27b_last_error()); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } std::printf("[target-split] gpu=%d layers=[%d,%d) %s\n", @@ -376,7 +362,7 @@ static int run_target_layer_split_harness( allocate_target_feat)) { std::fprintf(stderr, "target-split cache gpu=%d: %s\n", shard.gpu, dflash27b_last_error()); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } } @@ -395,7 +381,7 @@ static int run_target_layer_split_harness( if (!remote_draft.start(draft_ipc_bin, draft_path, draft_ipc_gpu, cap, draft_ipc_work_dir ? draft_ipc_work_dir : "")) { std::fprintf(stderr, "target-split remote draft start failed\n"); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } } else { @@ -409,7 +395,7 @@ static int run_target_layer_split_harness( draft_backend = ggml_backend_cuda_init(draft_gpu); if (!draft_backend) { std::fprintf(stderr, "target-split draft cuda init failed for gpu %d\n", draft_gpu); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } draft_backend_owned = true; @@ -426,7 +412,7 @@ static int run_target_layer_split_harness( draft_gpu, dflash27b_last_error()); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } std::printf("[target-split] draft loaded on gpu=%d format=%s\n", @@ -450,7 +436,7 @@ static int run_target_layer_split_harness( draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } std::printf("[target-split] draft feature ring cap=%d gpu=%d\n", cap, draft_gpu); @@ -463,7 +449,7 @@ static int run_target_layer_split_harness( draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } if ((int)prompt.size() + n_gen + 1 > max_ctx) { @@ -472,7 +458,7 @@ static int run_target_layer_split_harness( draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } @@ -485,7 +471,7 @@ static int run_target_layer_split_harness( int last_tok = -1; auto t_pf0 = std::chrono::steady_clock::now(); - if (!run_target_layer_split_forward(shards, shards.front().weights, + if (!run_qwen35_layer_split_forward(shards, shards.front().weights, prompt, 0, ubatch, last_tok, g_kq_stride_pad, g_fa_window, (load_draft && !use_remote_draft) ? &feature_ring : nullptr, @@ -495,7 +481,7 @@ static int run_target_layer_split_harness( draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } auto t_pf1 = std::chrono::steady_clock::now(); @@ -517,7 +503,7 @@ static int run_target_layer_split_harness( draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } if (use_remote_draft) { @@ -531,7 +517,7 @@ static int run_target_layer_split_harness( draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } std::printf("[target-split] remote draft smoke ctx=%d q=%d time=%.3f ms\n", @@ -551,7 +537,7 @@ static int run_target_layer_split_harness( draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } if (!use_mirror_view && @@ -563,7 +549,7 @@ static int run_target_layer_split_harness( draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } ggml_backend_tensor_set(draft_sg.inp_embed, noise_embed.data(), 0, @@ -584,7 +570,7 @@ static int run_target_layer_split_harness( draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } std::printf("[target-split] draft smoke ctx=%d q=%d time=%.3f ms\n", @@ -606,7 +592,7 @@ static int run_target_layer_split_harness( draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return ok ? 0 : 1; } @@ -616,7 +602,7 @@ static int run_target_layer_split_harness( for (; generated < n_gen; generated++) { std::vector one(1, last_tok); int next_tok = -1; - if (!run_target_layer_split_forward(shards, shards.front().weights, + if (!run_qwen35_layer_split_forward(shards, shards.front().weights, one, (int)out_all.size(), 1, next_tok, g_kq_stride_pad, g_fa_window, (load_draft && !use_remote_draft) ? &feature_ring : nullptr, @@ -626,7 +612,7 @@ static int run_target_layer_split_harness( draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 1; } out_all.push_back(last_tok); @@ -644,7 +630,7 @@ static int run_target_layer_split_harness( draft_feature_mirror_free(feature_ring); free_draft_weights(draft_weights); if (draft_backend_owned) ggml_backend_free(draft_backend); - free_target_layer_split_shards(shards); + free_qwen35_layer_split_shards(shards); return 0; } From 8415f1e9a208bddc957bb348d04e1d963547d1ff Mon Sep 17 00:00:00 2001 From: weicj Date: Sun, 24 May 2026 05:29:20 +0800 Subject: [PATCH 3/7] test(server): keep server unit tests CPU-only --- server/CMakeLists.txt | 16 +++--- server/src/common/layer_split_utils.cpp | 58 -------------------- server/src/common/layer_split_utils.h | 8 --- server/src/placement/placement_config.cpp | 65 +++++++++++++++++++++++ server/src/placement/placement_config.h | 8 +++ 5 files changed, 82 insertions(+), 73 deletions(-) create mode 100644 server/src/placement/placement_config.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 4be4744a4..6d874021a 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -266,6 +266,7 @@ add_library(dflash_common STATIC src/common/daemon_loop.cpp src/common/gguf_inspect.cpp src/common/backend_factory.cpp + src/placement/placement_config.cpp src/common/layer_split_utils.cpp src/common/ddtree.cpp src/common/peer_access.cpp @@ -772,6 +773,9 @@ if(DFLASH27B_TESTS) src/server/http_server.cpp src/server/model_card.cpp) target_include_directories(test_server_unit PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_include_directories(test_server_unit PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/common) if(DFLASH27B_GPU_BACKEND STREQUAL "hip") target_compile_definitions(test_server_unit PRIVATE DFLASH27B_BACKEND_HIP=1 GGML_USE_HIP) else() @@ -779,13 +783,11 @@ if(DFLASH27B_TESTS) DFLASH27B_BACKEND_CUDA=1 DFLASH27B_CUDA_MIN_SM=${_dflash_cuda_min_sm}) endif() - target_link_libraries(test_server_unit PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) - if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") - find_package(CUDAToolkit REQUIRED) - target_link_libraries(test_server_unit PRIVATE CUDA::cudart) - else() - target_link_libraries(test_server_unit PRIVATE hip::host) - endif() + target_link_libraries(test_server_unit PRIVATE + ggml-base + ggml-cpu + nlohmann_json::nlohmann_json + pthread) add_test(NAME server_unit COMMAND test_server_unit) endif() diff --git a/server/src/common/layer_split_utils.cpp b/server/src/common/layer_split_utils.cpp index e1f9f3f42..1f0cbfcf5 100644 --- a/server/src/common/layer_split_utils.cpp +++ b/server/src/common/layer_split_utils.cpp @@ -7,7 +7,6 @@ #include #include #include -#include namespace dflash::common { @@ -110,61 +109,4 @@ void free_layer_split_snapshot_backends( snapshot_backends.clear(); } -std::string validate_device_placement( - const DevicePlacement & dp, - int device_count) -{ - const bool validate_device_count = device_count >= 0; - if (validate_device_count && device_count == 0) { - return "no GPU devices available"; - } - - if (dp.gpu < 0 || - (validate_device_count && dp.gpu >= device_count)) { - return "primary gpu " + std::to_string(dp.gpu) + " out of range" + - (validate_device_count - ? " [0, " + std::to_string(device_count) + ")" - : ""); - } - - if (!dp.layer_split_gpus.empty()) { - if (dp.layer_split_gpus.size() < 2) { - return "layer_split_gpus must have at least 2 entries"; - } - - std::set seen; - for (int g : dp.layer_split_gpus) { - if (g < 0 || - (validate_device_count && g >= device_count)) { - return "layer_split gpu " + std::to_string(g) + - " out of range" + - (validate_device_count - ? " [0, " + std::to_string(device_count) + ")" - : ""); - } - if (!seen.insert(g).second) { - return "duplicate gpu " + std::to_string(g) + " in layer_split_gpus"; - } - } - - if (!dp.layer_split_weights.empty() && - dp.layer_split_weights.size() != dp.layer_split_gpus.size()) { - return "layer_split_weights size (" + - std::to_string(dp.layer_split_weights.size()) + - ") != gpu count (" + - std::to_string(dp.layer_split_gpus.size()) + ")"; - } - - for (double w : dp.layer_split_weights) { - if (w <= 0.0 || !std::isfinite(w)) { - return "layer_split_weights must be positive finite values"; - } - } - } - - if (dp.max_ctx <= 0) return "max_ctx must be positive"; - - return {}; // ok -} - } // namespace dflash::common diff --git a/server/src/common/layer_split_utils.h b/server/src/common/layer_split_utils.h index 447c41ccc..fb2facc0b 100644 --- a/server/src/common/layer_split_utils.h +++ b/server/src/common/layer_split_utils.h @@ -85,12 +85,4 @@ void free_layer_split_snapshot_backends( const std::vector & shards, std::vector & snapshot_backends); -// Validate a DevicePlacement against system constraints. -// If device_count is negative, only validates structural constraints that do -// not require querying the runtime-visible GPU count. -// Returns empty string on success, error description on failure. -std::string validate_device_placement( - const DevicePlacement & dp, - int device_count); - } // namespace dflash::common diff --git a/server/src/placement/placement_config.cpp b/server/src/placement/placement_config.cpp new file mode 100644 index 000000000..f7dcbd01f --- /dev/null +++ b/server/src/placement/placement_config.cpp @@ -0,0 +1,65 @@ +#include "placement_config.h" + +#include +#include + +namespace dflash::common { + +std::string validate_device_placement( + const DevicePlacement & dp, + int device_count) +{ + const bool validate_device_count = device_count >= 0; + if (validate_device_count && device_count == 0) { + return "no GPU devices available"; + } + + if (dp.gpu < 0 || + (validate_device_count && dp.gpu >= device_count)) { + return "primary gpu " + std::to_string(dp.gpu) + " out of range" + + (validate_device_count + ? " [0, " + std::to_string(device_count) + ")" + : ""); + } + + if (!dp.layer_split_gpus.empty()) { + if (dp.layer_split_gpus.size() < 2) { + return "layer_split_gpus must have at least 2 entries"; + } + + std::set seen; + for (int g : dp.layer_split_gpus) { + if (g < 0 || + (validate_device_count && g >= device_count)) { + return "layer_split gpu " + std::to_string(g) + + " out of range" + + (validate_device_count + ? " [0, " + std::to_string(device_count) + ")" + : ""); + } + if (!seen.insert(g).second) { + return "duplicate gpu " + std::to_string(g) + " in layer_split_gpus"; + } + } + + if (!dp.layer_split_weights.empty() && + dp.layer_split_weights.size() != dp.layer_split_gpus.size()) { + return "layer_split_weights size (" + + std::to_string(dp.layer_split_weights.size()) + + ") != gpu count (" + + std::to_string(dp.layer_split_gpus.size()) + ")"; + } + + for (double w : dp.layer_split_weights) { + if (w <= 0.0 || !std::isfinite(w)) { + return "layer_split_weights must be positive finite values"; + } + } + } + + if (dp.max_ctx <= 0) return "max_ctx must be positive"; + + return {}; +} + +} // namespace dflash::common diff --git a/server/src/placement/placement_config.h b/server/src/placement/placement_config.h index 63b5cf20b..3529065b0 100644 --- a/server/src/placement/placement_config.h +++ b/server/src/placement/placement_config.h @@ -102,4 +102,12 @@ inline bool parse_placement_device_list(const std::string & value, return true; } +// Validate a DevicePlacement against system constraints. +// If device_count is negative, only validates structural constraints that do +// not require querying the runtime-visible GPU count. +// Returns empty string on success, error description on failure. +std::string validate_device_placement( + const DevicePlacement & dp, + int device_count); + } // namespace dflash::common From b4ce59b1cb33f0d2314f77f2f7c2f34f9ce05679 Mon Sep 17 00:00:00 2001 From: weicj Date: Sun, 24 May 2026 15:04:34 +0800 Subject: [PATCH 4/7] fix(server): restore DFlash state for split targets --- server/src/common/dflash_draft_ipc.cpp | 58 +++++++++ server/src/common/dflash_draft_ipc.h | 4 + server/src/common/dflash_draft_ipc_daemon.cpp | 63 ++++++++++ server/src/common/layer_split_backend.cpp | 2 +- .../src/qwen35/qwen35_layer_split_adapter.cpp | 116 ++++++++++++++++++ .../src/qwen35/qwen35_layer_split_adapter.h | 13 ++ server/src/server/http_server.cpp | 3 +- server/test/test_server_unit.cpp | 23 +++- 8 files changed, 276 insertions(+), 6 deletions(-) diff --git a/server/src/common/dflash_draft_ipc.cpp b/server/src/common/dflash_draft_ipc.cpp index 0f8542379..95045e97a 100644 --- a/server/src/common/dflash_draft_ipc.cpp +++ b/server/src/common/dflash_draft_ipc.cpp @@ -122,6 +122,64 @@ bool DFlashDraftIpcClient::propose( #endif } +bool DFlashDraftIpcClient::get_feature_range(int start_pos, int n_tokens, + std::vector & out) { +#if defined(_WIN32) + (void)start_pos; (void)n_tokens; (void)out; + return false; +#else + FILE * cmd = process_.command_stream(); + const int stream_fd = process_.stream_fd(); + if (!active_ || !cmd || stream_fd < 0 || ring_cap_ <= 0 || + start_pos < 0 || n_tokens <= 0 || n_tokens > ring_cap_) return false; + const size_t count = + (size_t)n_tokens * (size_t)n_target_layers_ * (size_t)hidden_size_; + std::fprintf(cmd, "get_feature_range %d %d\n", start_pos, n_tokens); + std::fflush(cmd); + int32_t status = -1; + bool ok = read_exact_fd(stream_fd, &status, sizeof(status)) && status == 0; + if (ok) { + out.assign(count, 0.0f); + ok = read_exact_fd(stream_fd, out.data(), out.size() * sizeof(float)); + } + if (!ok) { + std::fprintf(stderr, "draft-ipc get_feature_range failed status=%d\n", status); + } + return ok; +#endif +} + +bool DFlashDraftIpcClient::set_feature_range(int start_pos, int n_tokens, + const std::vector & data) { +#if defined(_WIN32) + (void)start_pos; (void)n_tokens; (void)data; + return false; +#else + FILE * cmd = process_.command_stream(); + const int stream_fd = process_.stream_fd(); + if (!active_ || !cmd || stream_fd < 0 || ring_cap_ <= 0 || + start_pos < 0 || n_tokens <= 0 || n_tokens > ring_cap_) return false; + const size_t expected = + (size_t)n_tokens * (size_t)n_target_layers_ * (size_t)hidden_size_; + if (data.size() != expected) return false; + const std::string path = process_.next_path("feature_range"); + if (!write_binary_file(path, data.data(), data.size() * sizeof(float))) { + std::fprintf(stderr, "draft-ipc write feature range failed: %s\n", path.c_str()); + return false; + } + std::fprintf(cmd, "set_feature_range %d %d %s\n", + start_pos, n_tokens, path.c_str()); + std::fflush(cmd); + int32_t status = -1; + const bool ok = read_exact_fd(stream_fd, &status, sizeof(status)) && status == 0; + std::remove(path.c_str()); + if (!ok) { + std::fprintf(stderr, "draft-ipc set_feature_range failed status=%d\n", status); + } + return ok; +#endif +} + void DFlashDraftIpcClient::close() { process_.close(); active_ = false; diff --git a/server/src/common/dflash_draft_ipc.h b/server/src/common/dflash_draft_ipc.h index f719555ec..3a361b8b0 100644 --- a/server/src/common/dflash_draft_ipc.h +++ b/server/src/common/dflash_draft_ipc.h @@ -55,6 +55,10 @@ class DFlashDraftIpcClient { const std::vector & noise_embed, std::vector & hidden_out); + bool get_feature_range(int start_pos, int n_tokens, std::vector & out); + bool set_feature_range(int start_pos, int n_tokens, + const std::vector & data); + bool active() const { return active_; } int ring_cap() const { return ring_cap_; } int hidden_size() const { return hidden_size_; } diff --git a/server/src/common/dflash_draft_ipc_daemon.cpp b/server/src/common/dflash_draft_ipc_daemon.cpp index 9cbc019f7..953e0abb7 100644 --- a/server/src/common/dflash_draft_ipc_daemon.cpp +++ b/server/src/common/dflash_draft_ipc_daemon.cpp @@ -132,6 +132,69 @@ int run_dflash_draft_ipc_daemon(const char * draft_path, stream_status(stream_fd, 0); continue; } + if (cmd == "get_feature_range") { + int start_pos = -1; + int n_tokens = 0; + iss >> start_pos >> n_tokens; + if (start_pos < 0 || n_tokens <= 0 || n_tokens > feature_ring.cap) { + std::fprintf(stderr, "[draft-ipc-daemon] bad get_feature_range: %s\n", + line.c_str()); + stream_status(stream_fd, -1); + continue; + } + const int fc_in = feature_ring.n_target_layers * feature_ring.hidden_size; + const size_t row_bytes = (size_t)fc_in * sizeof(float); + const size_t src_stride = feature_ring.target_feat->nb[1]; + std::vector data((size_t)n_tokens * (size_t)fc_in); + for (int i = 0; i < n_tokens; ++i) { + const int slot = (start_pos + i) % feature_ring.cap; + ggml_backend_tensor_get(feature_ring.target_feat, + data.data() + (size_t)i * (size_t)fc_in, + (size_t)slot * src_stride, + row_bytes); + } + const size_t bytes = data.size() * sizeof(float); + if (!stream_status(stream_fd, 0) || + !write_exact_fd(stream_fd, data.data(), bytes)) { + std::fprintf(stderr, "[draft-ipc-daemon] feature range stream failed\n"); + break; + } + continue; + } + if (cmd == "set_feature_range") { + int start_pos = -1; + int n_tokens = 0; + iss >> start_pos >> n_tokens; + std::string path = read_line_tail(iss); + if (start_pos < 0 || n_tokens <= 0 || n_tokens > feature_ring.cap || + path.empty()) { + std::fprintf(stderr, "[draft-ipc-daemon] bad set_feature_range: %s\n", + line.c_str()); + stream_status(stream_fd, -1); + continue; + } + const int fc_in = feature_ring.n_target_layers * feature_ring.hidden_size; + const size_t row_bytes = (size_t)fc_in * sizeof(float); + const size_t dst_stride = feature_ring.target_feat->nb[1]; + const size_t bytes = (size_t)n_tokens * row_bytes; + std::vector data(bytes / sizeof(float)); + if (!read_binary_file_exact(path, data.data(), bytes)) { + std::fprintf(stderr, "[draft-ipc-daemon] read feature range failed: %s\n", + path.c_str()); + stream_status(stream_fd, -1); + continue; + } + for (int i = 0; i < n_tokens; ++i) { + const int slot = (start_pos + i) % feature_ring.cap; + ggml_backend_tensor_set(feature_ring.target_feat, + data.data() + (size_t)i * (size_t)fc_in, + (size_t)slot * dst_stride, + row_bytes); + } + ggml_backend_synchronize(backend); + stream_status(stream_fd, 0); + continue; + } if (cmd == "propose") { int committed = -1; int ctx_len = 0; diff --git a/server/src/common/layer_split_backend.cpp b/server/src/common/layer_split_backend.cpp index 13cbc2c87..18ef3980d 100644 --- a/server/src/common/layer_split_backend.cpp +++ b/server/src/common/layer_split_backend.cpp @@ -101,7 +101,7 @@ GenerateResult LayerSplitBackend::run_from_state(const GenerateRequest & req, return result; } auto t_decode_start = std::chrono::steady_clock::now(); - const bool ok = (base_pos == 0 && adapter_->can_dflash_decode()) + const bool ok = adapter_->can_dflash_decode() ? adapter_->decode_dflash(req.prompt, base_pos, last_tok, req.n_gen, result.tokens, out_io) : adapter_->decode_ar(last_tok, base_pos + (int)req.prompt.size(), req.n_gen, diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index 87e6d067c..a39f8bacd 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -86,6 +86,7 @@ bool Qwen35LayerSplitAdapter::init() { for (auto & slot : prefix_snapshots_) { slot.resize(shards_.size()); } + draft_feature_snapshots_.resize(PREFIX_SLOTS); return true; } @@ -208,6 +209,10 @@ bool Qwen35LayerSplitAdapter::snapshot_save(int slot) { return false; } } + if (!snapshot_draft_features(slot)) { + snapshot_free(slot); + return false; + } return true; } @@ -216,6 +221,7 @@ void Qwen35LayerSplitAdapter::snapshot_free(int slot) { for (auto & snap : prefix_snapshots_[(size_t)slot]) { free_prefix_snapshot(snap); } + free_draft_feature_snapshot(slot); } bool Qwen35LayerSplitAdapter::snapshot_used(int slot) const { @@ -225,6 +231,12 @@ bool Qwen35LayerSplitAdapter::snapshot_used(int slot) const { for (const auto & snap : snaps) { if (!snap.ctx) return false; } + if (cfg_.run_dflash && cfg_.draft_path) { + if (draft_feature_snapshots_.size() != (size_t)PREFIX_SLOTS) return false; + const auto & draft_snap = draft_feature_snapshots_[(size_t)slot]; + if (draft_snap.cur_pos <= 0 || draft_snap.n_tokens <= 0 || + draft_snap.data.empty()) return false; + } return true; } @@ -243,6 +255,109 @@ bool Qwen35LayerSplitAdapter::snapshot_restore(int slot) { return false; } } + if (!restore_draft_features(slot)) return false; + return true; +} + +bool Qwen35LayerSplitAdapter::snapshot_draft_features(int slot) { + if (!cfg_.run_dflash || !cfg_.draft_path) { + free_draft_feature_snapshot(slot); + return true; + } + if (!snapshot_slot_valid(slot) || + draft_feature_snapshots_.size() != (size_t)PREFIX_SLOTS) { + return false; + } + + const auto & snaps = prefix_snapshots_[(size_t)slot]; + if (snaps.empty() || !snaps.front().ctx) return false; + const int cur_pos = snaps.front().cur_pos; + if (cur_pos <= 0) return false; + const int ring_cap = remote_draft_.active() ? remote_draft_.ring_cap() : feature_ring_.cap; + const int n_layers = remote_draft_.active() ? remote_draft_.n_target_layers() + : feature_ring_.n_target_layers; + const int hidden = remote_draft_.active() ? remote_draft_.hidden_size() + : feature_ring_.hidden_size; + if (ring_cap <= 0 || n_layers <= 0 || hidden <= 0) return false; + const int n_tokens = std::min(cur_pos, ring_cap); + const int start_pos = cur_pos - n_tokens; + if (n_tokens <= 0) return false; + + auto & snap = draft_feature_snapshots_[(size_t)slot]; + snap.cur_pos = cur_pos; + snap.start_pos = start_pos; + snap.n_tokens = n_tokens; + snap.cap = ring_cap; + snap.n_target_layers = n_layers; + snap.hidden_size = hidden; + snap.data.clear(); + snap.data.resize((size_t)n_tokens * (size_t)n_layers * (size_t)hidden); + + if (remote_draft_.active()) { + return remote_draft_.get_feature_range(start_pos, n_tokens, snap.data); + } + + if (!feature_ring_.target_feat) return false; + const int fc_in = n_layers * hidden; + const size_t row_bytes = (size_t)fc_in * sizeof(float); + const size_t src_stride = feature_ring_.target_feat->nb[1]; + for (int i = 0; i < n_tokens; ++i) { + const int ring_slot = (start_pos + i) % ring_cap; + ggml_backend_tensor_get(feature_ring_.target_feat, + snap.data.data() + (size_t)i * (size_t)fc_in, + (size_t)ring_slot * src_stride, + row_bytes); + } + return true; +} + +void Qwen35LayerSplitAdapter::free_draft_feature_snapshot(int slot) { + if (slot < 0 || draft_feature_snapshots_.size() != (size_t)PREFIX_SLOTS || + slot >= (int)draft_feature_snapshots_.size()) { + return; + } + draft_feature_snapshots_[(size_t)slot] = DraftFeatureSnapshot{}; +} + +bool Qwen35LayerSplitAdapter::restore_draft_features(int slot) { + if (!cfg_.run_dflash || !cfg_.draft_path) return true; + if (slot < 0 || draft_feature_snapshots_.size() != (size_t)PREFIX_SLOTS || + slot >= (int)draft_feature_snapshots_.size()) { + return false; + } + + const auto & snap = draft_feature_snapshots_[(size_t)slot]; + if (snap.cur_pos <= 0 || snap.start_pos < 0 || snap.n_tokens <= 0 || + snap.cap <= 0 || snap.n_target_layers <= 0 || snap.hidden_size <= 0 || + snap.data.empty()) { + return false; + } + + if (remote_draft_.active()) { + if (snap.cap != remote_draft_.ring_cap() || + snap.n_target_layers != remote_draft_.n_target_layers() || + snap.hidden_size != remote_draft_.hidden_size()) { + return false; + } + return remote_draft_.set_feature_range(snap.start_pos, snap.n_tokens, snap.data); + } + + if (!feature_ring_.target_feat || + snap.cap != feature_ring_.cap || + snap.n_target_layers != feature_ring_.n_target_layers || + snap.hidden_size != feature_ring_.hidden_size) { + return false; + } + const int fc_in = snap.n_target_layers * snap.hidden_size; + const size_t row_bytes = (size_t)fc_in * sizeof(float); + const size_t dst_stride = feature_ring_.target_feat->nb[1]; + for (int i = 0; i < snap.n_tokens; ++i) { + const int ring_slot = (snap.start_pos + i) % snap.cap; + ggml_backend_tensor_set(feature_ring_.target_feat, + snap.data.data() + (size_t)i * (size_t)fc_in, + (size_t)ring_slot * dst_stride, + row_bytes); + } return true; } @@ -378,6 +493,7 @@ void Qwen35LayerSplitAdapter::shutdown() { for (auto & snap : slot) free_prefix_snapshot(snap); } prefix_snapshots_.clear(); + draft_feature_snapshots_.clear(); auto shard_metas = layer_split_shard_metas(shards_); free_layer_split_snapshot_backends(shard_metas, snapshot_backends_); if (draft_backend_owned_ && draft_backend_) { diff --git a/server/src/qwen35/qwen35_layer_split_adapter.h b/server/src/qwen35/qwen35_layer_split_adapter.h index 1feaf82dd..b4ab49662 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.h +++ b/server/src/qwen35/qwen35_layer_split_adapter.h @@ -81,6 +81,9 @@ class Qwen35LayerSplitAdapter : public LayerSplitAdapter { private: bool load_draft(); bool snapshot_slot_valid(int slot) const; + bool snapshot_draft_features(int slot); + void free_draft_feature_snapshot(int slot); + bool restore_draft_features(int slot); Qwen35LayerSplitAdapterConfig cfg_; std::vector shards_; @@ -96,6 +99,16 @@ class Qwen35LayerSplitAdapter : public LayerSplitAdapter { static constexpr int PREFIX_SLOTS = ModelBackend::kMaxSlots; std::vector> prefix_snapshots_; std::vector snapshot_backends_; + struct DraftFeatureSnapshot { + int cur_pos = 0; + int start_pos = 0; + int n_tokens = 0; + int cap = 0; + int n_target_layers = 0; + int hidden_size = 0; + std::vector data; + }; + std::vector draft_feature_snapshots_; SamplerCfg sampler_; std::mt19937_64 sampler_rng_{std::random_device{}()}; diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 0f175ec69..ab37805bf 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -1508,7 +1508,8 @@ void HttpServer::worker_loop() { // Confirm or abort the inline snapshot. if (snap_prepared) { - if (completion_tokens > 0 && !client_disconnected) { + if (completion_tokens > 0 && !client_disconnected && + backend_.snapshot_used(snap_slot)) { prefix_cache_.confirm_inline_snap(snap_slot, snap_cut, effective_prompt); // Track for shutdown save. slot_tokens_[snap_slot] = std::vector( diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index c3b803a6a..eddc7741e 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -1241,6 +1241,8 @@ struct MockLayerSplitAdapter : LayerSplitAdapter { int current_last = -1; std::vector prefill_bases; std::vector prefill_sizes; + int dflash_base = -1; + int dflash_last = -1; std::vector emitted_tokens; bool dflash_enabled = false; bool dflash_called = false; @@ -1277,10 +1279,21 @@ struct MockLayerSplitAdapter : LayerSplitAdapter { return true; } bool can_dflash_decode() const override { return dflash_enabled; } - bool decode_dflash(const std::vector &, int, int, int, - std::vector &, const DaemonIO &) override { + bool decode_dflash(const std::vector & prompt, int base_pos, + int last_tok, int n_gen, std::vector & out_tokens, + const DaemonIO & io) override { + (void)prompt; dflash_called = true; - return false; + dflash_base = base_pos; + dflash_last = last_tok; + for (int i = 0; i < n_gen; ++i) { + int32_t tok = last_tok + i + 10; + out_tokens.push_back(tok); + emitted_tokens.push_back(tok); + io.emit(tok); + } + io.emit(-1); + return true; } void free_drafter() override {} bool snapshot_save(int slot) override { @@ -1350,12 +1363,14 @@ static void test_layer_split_backend_inline_snapshot_and_restore_delta() { GenerateResult restored = backend.restore_and_generate(2, restore_req, io); TEST_ASSERT(restored.ok); - TEST_ASSERT(!raw->dflash_called); + TEST_ASSERT(raw->dflash_called); TEST_ASSERT(raw->restored_slot == 2); TEST_ASSERT(!raw->reset_called); TEST_ASSERT(raw->prefill_bases.size() == 1); TEST_ASSERT(raw->prefill_bases[0] == 3); TEST_ASSERT(raw->prefill_sizes[0] == 1); + TEST_ASSERT(raw->dflash_base == 3); + TEST_ASSERT(raw->dflash_last == 99); } static void test_layer_split_compress_nopark_uses_default_drafter_path() { From 061e31ce2c309d40e353ba25aab8d361d8c508fe Mon Sep 17 00:00:00 2001 From: weicj Date: Thu, 28 May 2026 13:57:41 +0800 Subject: [PATCH 5/7] fix(server): align layer split load plan with target load plan --- server/src/common/layer_split_utils.h | 5 +++-- server/src/internal.h | 1 - server/src/qwen35/layer_split_daemon_loop.cpp | 2 +- server/src/qwen35/qwen35_layer_split_adapter.cpp | 2 +- server/test/test_dflash.cpp | 2 +- 5 files changed, 6 insertions(+), 6 deletions(-) diff --git a/server/src/common/layer_split_utils.h b/server/src/common/layer_split_utils.h index fb2facc0b..ae9029526 100644 --- a/server/src/common/layer_split_utils.h +++ b/server/src/common/layer_split_utils.h @@ -26,10 +26,11 @@ struct LayerSplitShardMeta { ggml_backend_t backend = nullptr; }; -inline TargetLoadPlan make_layer_split_load_plan( +template +inline LoadPlan make_layer_split_load_plan( const LayerSplitShardMeta & shard, bool is_last_shard) { - TargetLoadPlan plan; + LoadPlan plan; plan.layer_begin = shard.layer_begin; plan.layer_end = shard.layer_end; plan.load_output = is_last_shard; diff --git a/server/src/internal.h b/server/src/internal.h index adeab6321..f9a890ff2 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -23,7 +23,6 @@ #include "ggml-backend.h" #include "gguf.h" -#include "common/layer_split_utils.h" #include "dflash27b.h" namespace dflash::common { diff --git a/server/src/qwen35/layer_split_daemon_loop.cpp b/server/src/qwen35/layer_split_daemon_loop.cpp index be2d4d535..fe2862756 100644 --- a/server/src/qwen35/layer_split_daemon_loop.cpp +++ b/server/src/qwen35/layer_split_daemon_loop.cpp @@ -49,7 +49,7 @@ int run_layer_split_daemon(const LayerSplitDaemonConfig & cfg) { // Load partial target weights + caches. for (auto & shard : shards) { const TargetLoadPlan plan = - make_layer_split_load_plan(shard, &shard == &shards.back()); + make_layer_split_load_plan(shard, &shard == &shards.back()); if (!load_target_gguf_partial(cfg.target_path, shard.backend, plan, shard.weights) || !create_target_cache_partial(shard.weights, cfg.max_ctx, cfg.max_verify_tokens, shard.backend, shard.cache, diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index a39f8bacd..918eed5ea 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -62,7 +62,7 @@ bool Qwen35LayerSplitAdapter::init() { for (auto & shard : shards_) { const TargetLoadPlan plan = - make_layer_split_load_plan(shard, &shard == &shards_.back()); + make_layer_split_load_plan(shard, &shard == &shards_.back()); if (!load_target_gguf_partial(cfg_.target_path, shard.backend, plan, shard.weights) || !create_target_cache_partial(shard.weights, cfg_.device.max_ctx, diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index 26e51e4c2..93bd1b1d2 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -344,7 +344,7 @@ static int run_target_layer_split_harness( (void)enable_layer_split_peer_access(target_gpus, peer_access); for (auto & shard : shards) { const TargetLoadPlan plan = - make_layer_split_load_plan(shard, &shard == &shards.back()); + make_layer_split_load_plan(shard, &shard == &shards.back()); if (!load_target_gguf_partial(target_path, shard.backend, plan, shard.weights)) { std::fprintf(stderr, "target-split load gpu=%d: %s\n", shard.gpu, dflash27b_last_error()); From 73c4a85c9c4e93ad315593f220f1c72f9cd767a1 Mon Sep 17 00:00:00 2001 From: weicj Date: Thu, 28 May 2026 14:06:58 +0800 Subject: [PATCH 6/7] fix(server): link layer split unit tests with common library --- server/CMakeLists.txt | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 6d874021a..d6f2b0db6 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -773,9 +773,6 @@ if(DFLASH27B_TESTS) src/server/http_server.cpp src/server/model_card.cpp) target_include_directories(test_server_unit PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) - target_include_directories(test_server_unit PRIVATE - ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src - ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/common) if(DFLASH27B_GPU_BACKEND STREQUAL "hip") target_compile_definitions(test_server_unit PRIVATE DFLASH27B_BACKEND_HIP=1 GGML_USE_HIP) else() @@ -783,11 +780,13 @@ if(DFLASH27B_TESTS) DFLASH27B_BACKEND_CUDA=1 DFLASH27B_CUDA_MIN_SM=${_dflash_cuda_min_sm}) endif() - target_link_libraries(test_server_unit PRIVATE - ggml-base - ggml-cpu - nlohmann_json::nlohmann_json - pthread) + target_link_libraries(test_server_unit PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + find_package(CUDAToolkit REQUIRED) + target_link_libraries(test_server_unit PRIVATE CUDA::cudart) + else() + target_link_libraries(test_server_unit PRIVATE hip::host) + endif() add_test(NAME server_unit COMMAND test_server_unit) endif() From 054af280a6283e20144ebd288ab6fb834725d6a4 Mon Sep 17 00:00:00 2001 From: weicj Date: Thu, 28 May 2026 15:50:21 +0800 Subject: [PATCH 7/7] fix(server): tighten layer-split validation and cleanup --- server/src/common/backend_factory.cpp | 1 + server/src/common/layer_split_backend.cpp | 12 ++++++++ server/src/common/layer_split_backend.h | 1 + .../src/qwen35/qwen35_layer_split_adapter.cpp | 7 +++++ .../src/qwen35/qwen35_layer_split_adapter.h | 1 + server/src/server/server_main.cpp | 9 ++++-- server/test/test_server_unit.cpp | 29 ++++++++++++++++++- 7 files changed, 56 insertions(+), 4 deletions(-) diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 2075fd1cc..977a7c762 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -54,6 +54,7 @@ std::unique_ptr create_backend(const BackendArgs & args) { cfg.remote_draft = args.remote_draft; cfg.fa_window = args.fa_window; cfg.kq_stride_pad = args.kq_stride_pad; + cfg.draft_swa_window = args.draft_swa_window; cfg.draft_ctx_max = args.draft_ctx_max; cfg.max_verify_tokens = args.ddtree_mode ? std::max(DFLASH27B_DRAFT_BLOCK_SIZE, args.ddtree_budget + 1) diff --git a/server/src/common/layer_split_backend.cpp b/server/src/common/layer_split_backend.cpp index 18ef3980d..11e75e0a7 100644 --- a/server/src/common/layer_split_backend.cpp +++ b/server/src/common/layer_split_backend.cpp @@ -6,6 +6,7 @@ #include #include +#include #include namespace dflash::common { @@ -20,6 +21,7 @@ bool LayerSplitBackend::init() { std::fprintf(stderr, "[target-split] missing model adapter\n"); return false; } + shutdown_done_ = false; return adapter_->init(); } @@ -189,6 +191,14 @@ bool LayerSplitBackend::handle_compress(const std::string & line, CompressRequest req; req.input_ids = read_int32_file(ppath); req.keep_ratio = (float)keep_x1000 / 1000.0f; + if (!std::isfinite(req.keep_ratio) || + req.keep_ratio < 0.0f || req.keep_ratio > 1.0f) { + std::fprintf(stderr, + "[target-split][compress] keep ratio must be in [0,1], got %d/1000\n", + keep_x1000); + io.emit(-1); + return false; + } if (n >= 3 && drafter_path[0]) { req.drafter_path = drafter_path; } else if (adapter_) { @@ -219,6 +229,8 @@ bool LayerSplitBackend::supports_remote_draft() const { } void LayerSplitBackend::shutdown() { + if (shutdown_done_) return; + shutdown_done_ = true; if (adapter_) adapter_->shutdown(); } diff --git a/server/src/common/layer_split_backend.h b/server/src/common/layer_split_backend.h index 4553b42e3..4aeda23f6 100644 --- a/server/src/common/layer_split_backend.h +++ b/server/src/common/layer_split_backend.h @@ -107,6 +107,7 @@ class LayerSplitBackend : public ModelBackend { bool reset_state); std::unique_ptr adapter_; + bool shutdown_done_ = false; }; } // namespace dflash::common diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index 918eed5ea..a9ece2c37 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -107,6 +107,9 @@ bool Qwen35LayerSplitAdapter::load_draft() { draft_weights_.n_embd = DFLASH27B_TARGET_HIDDEN; draft_weights_.block_size = DFLASH27B_DRAFT_BLOCK_SIZE; draft_weights_.n_target_layers = DFLASH27B_DRAFT_N_TARGET_LAYERS; + if (cfg_.draft_swa_window > 0) { + draft_weights_.swa_window = cfg_.draft_swa_window; + } std::fprintf(stderr, "[target-split] remote draft ready gpu=%d cap=%d\n", cfg_.draft_gpu, cap); @@ -142,6 +145,9 @@ bool Qwen35LayerSplitAdapter::load_draft() { cfg_.draft_gpu, dflash27b_last_error()); return false; } + if (cfg_.draft_swa_window > 0) { + draft_weights_.swa_window = cfg_.draft_swa_window; + } const int cap = std::min(cfg_.device.max_ctx, cfg_.draft_ctx_max); if (!draft_feature_mirror_init(feature_ring_, draft_backend_, @@ -502,6 +508,7 @@ void Qwen35LayerSplitAdapter::shutdown() { draft_backend_ = nullptr; draft_backend_owned_ = false; free_qwen35_layer_split_shards(shards_); + shards_.clear(); } } // namespace dflash::common diff --git a/server/src/qwen35/qwen35_layer_split_adapter.h b/server/src/qwen35/qwen35_layer_split_adapter.h index b4ab49662..4011d837a 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.h +++ b/server/src/qwen35/qwen35_layer_split_adapter.h @@ -33,6 +33,7 @@ struct Qwen35LayerSplitAdapterConfig { int draft_ctx_max = 4096; int max_verify_tokens = DFLASH27B_DRAFT_BLOCK_SIZE; bool run_dflash = false; + int draft_swa_window = 0; }; class Qwen35LayerSplitAdapter : public LayerSplitAdapter { diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index fc06e2c1b..3dcb23a5a 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -128,6 +128,12 @@ static bool validate_server_placement(const BackendArgs & bargs, placement_error.c_str()); return false; } + if (!sconfig.disk_cache_dir.empty()) { + std::fprintf(stderr, + "[server] --kv-cache-dir is not supported with --target-devices yet; " + "sharded disk snapshot/restore will be added separately\n"); + return false; + } } if (bargs.device.is_layer_split() && target != compiled) { std::fprintf(stderr, @@ -461,9 +467,6 @@ int main(int argc, char ** argv) { sconfig.pflash_drafter_gpu = pflash_placement.drafter_gpu; sconfig.pflash_remote_drafter = pflash_placement.remote_drafter; sconfig.pflash_remote = pflash_placement.remote; - if (bargs.device.is_layer_split()) { - sconfig.disk_cache_dir.clear(); - } // ── Apply environment defaults ───────────────────────────────────── // Explicit --cache-type-k/v override via env vars. diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index eddc7741e..1415aab30 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -1246,6 +1246,7 @@ struct MockLayerSplitAdapter : LayerSplitAdapter { std::vector emitted_tokens; bool dflash_enabled = false; bool dflash_called = false; + int shutdown_calls = 0; ModelBackend::CompressRequest last_compress_req; const char * name() const override { return "mock"; } @@ -1326,7 +1327,7 @@ struct MockLayerSplitAdapter : LayerSplitAdapter { result.compressed_ids = {77, 88}; return result; } - void shutdown() override {} + void shutdown() override { shutdown_calls++; } }; static void test_layer_split_backend_inline_snapshot_and_restore_delta() { @@ -1393,6 +1394,30 @@ static void test_layer_split_compress_nopark_uses_default_drafter_path() { unlink(ids_path.c_str()); } +static void test_layer_split_compress_rejects_bad_keep_ratio() { + const std::string ids_path = "/tmp/dflash_test_layer_split_compress_bad.bin"; + unlink(ids_path.c_str()); + TEST_ASSERT(write_int32_file(ids_path, {1, 2, 3, 4})); + + auto * raw = new MockLayerSplitAdapter(); + LayerSplitBackend backend{std::unique_ptr(raw)}; + DaemonIO io; + + const std::string cmd = "compress " + ids_path + " 1250 nopark"; + TEST_ASSERT(!backend.handle_compress(cmd, io)); + TEST_ASSERT(raw->last_compress_req.input_ids.empty()); + + unlink(ids_path.c_str()); +} + +static void test_layer_split_backend_shutdown_is_idempotent() { + auto * raw = new MockLayerSplitAdapter(); + LayerSplitBackend backend{std::unique_ptr(raw)}; + backend.shutdown(); + backend.shutdown(); + TEST_ASSERT(raw->shutdown_calls == 1); +} + // Disk Prefix Cache Tests // ═══════════════════════════════════════════════════════════════════════ @@ -2524,6 +2549,8 @@ int main() { RUN_TEST(test_validate_layer_split_weights_shape); RUN_TEST(test_layer_split_backend_inline_snapshot_and_restore_delta); RUN_TEST(test_layer_split_compress_nopark_uses_default_drafter_path); + RUN_TEST(test_layer_split_compress_rejects_bad_keep_ratio); + RUN_TEST(test_layer_split_backend_shutdown_is_idempotent); std::fprintf(stderr, "\n── Disk prefix cache ──\n"); RUN_TEST(test_disk_cache_config_defaults);