fix(server): restore DFlash state for split targets

weicj · weicj · commit b4ce59b1cb33 · 2026-05-28T13:53:13.000+08:00
diff --git 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<float> & 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<float> & 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
@@ -55,6 +55,10 @@ class DFlashDraftIpcClient {
                  const std::vector<float> & noise_embed,
                  std::vector<float> & hidden_out);
 
+    bool get_feature_range(int start_pos, int n_tokens, std::vector<float> & out);
+    bool set_feature_range(int start_pos, int n_tokens,
+                           const std::vector<float> & 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
@@ -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<float> 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<float> 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
@@ -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
@@ -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
@@ -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<Qwen35LayerSplitShard> shards_;
@@ -96,6 +99,16 @@ class Qwen35LayerSplitAdapter : public LayerSplitAdapter {
     static constexpr int PREFIX_SLOTS = ModelBackend::kMaxSlots;
     std::vector<std::vector<PrefixSnapshot>> prefix_snapshots_;
     std::vector<ggml_backend_t> 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<float> data;
+    };
+    std::vector<DraftFeatureSnapshot> 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
@@ -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<int32_t>(
diff --git 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<int> prefill_bases;
     std::vector<int> prefill_sizes;
+    int dflash_base = -1;
+    int dflash_last = -1;
     std::vector<int32_t> 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<int32_t> &, int, int, int,
-                       std::vector<int32_t> &, const DaemonIO &) override {
+    bool decode_dflash(const std::vector<int32_t> & prompt, int base_pos,
+                       int last_tok, int n_gen, std::vector<int32_t> & 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() {

Original file line number	Diff line number	Diff line change
`@@ -101,7 +101,7 @@ GenerateResult LayerSplitBackend::run_from_state(const GenerateRequest & req,`
`101`	`101`	`return result;`
`102`	`102`	`}`
`103`	`103`	`auto t_decode_start = std::chrono::steady_clock::now();`
`104`		`- const bool ok = (base_pos == 0 && adapter_->can_dflash_decode())`
	`104`	`+ const bool ok = adapter_->can_dflash_decode()`
`105`	`105`	`? adapter_->decode_dflash(req.prompt, base_pos, last_tok, req.n_gen,`
`106`	`106`	`result.tokens, out_io)`
`107`	`107`	`: adapter_->decode_ar(last_tok, base_pos + (int)req.prompt.size(), req.n_gen,`