refactor(server): generalize target layer-split adapter path

Author: weicj

dflash/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
 

dflash/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 <algorithm>
 #include <cmath>
+#include <cstdio>
 #include <set>
 
 namespace dflash::common {
 
-std::vector<std::pair<int,int>> compute_layer_ranges(
+std::vector<LayerSplitRange> compute_layer_ranges(
     int n_layer,
     int n_gpus,
     const std::vector<double> & weights)
 {
-    std::vector<std::pair<int,int>> ranges;
+    std::vector<LayerSplitRange> ranges;
     if (n_layer <= 0 || n_gpus <= 0 || n_gpus > n_layer) return ranges;
 
     std::vector<double> w = weights;
@@ -39,6 +44,72 @@ std::vector<std::pair<int,int>> compute_layer_ranges(
     return ranges;
 }
 
+bool init_layer_split_shard_metas(
+        std::vector<LayerSplitShardMeta *> shards,
+        const std::vector<int> & gpus,
+        const std::vector<LayerSplitRange> & 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<int> & 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<LayerSplitShardMeta *> & shards,
+        std::vector<ggml_backend_t> & 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<LayerSplitShardMeta *> & shards,
+        std::vector<ggml_backend_t> & 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)

dflash/src/common/layer_split_utils.h

@@ -8,20 +8,89 @@
 #include "placement/placement_config.h"
 
 #include <string>
-#include <utility>
 #include <vector>
 
+#include "ggml-backend.h"
+
 namespace dflash::common {
 
+struct LayerSplitRange {
+    int begin = 0;
+    int end = 0;
+};
+
+struct LayerSplitLoadPlan {
+    int  layer_begin = 0;     // inclusive
+    int  layer_end   = -1;    // exclusive; <0 means all layers
+    bool load_output = true;  // final output norm / lm-head tensors
+};
+
+struct LayerSplitShardMeta {
+    int gpu = 0;
+    int layer_begin = 0;
+    int layer_end = 0;
+    ggml_backend_t backend = nullptr;
+};
+
+inline LayerSplitLoadPlan make_layer_split_load_plan(
+        const LayerSplitShardMeta & shard,
+        bool is_last_shard) {
+    LayerSplitLoadPlan plan;
+    plan.layer_begin = shard.layer_begin;
+    plan.layer_end = shard.layer_end;
+    plan.load_output = is_last_shard;
+    return plan;
+}
+
+template <typename Shard>
+std::vector<LayerSplitShardMeta *> layer_split_shard_metas(
+        std::vector<Shard> & shards) {
+    std::vector<LayerSplitShardMeta *> metas;
+    metas.reserve(shards.size());
+    for (auto & shard : shards) {
+        metas.push_back(&shard);
+    }
+    return metas;
+}
+
+template <typename Shard>
+Shard * find_layer_split_shard(std::vector<Shard> & 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<std::pair<int,int>> compute_layer_ranges(
+std::vector<LayerSplitRange> compute_layer_ranges(
     int n_layer,
     int n_gpus,
     const std::vector<double> & weights);
 
+bool init_layer_split_shard_metas(
+    std::vector<LayerSplitShardMeta *> shards,
+    const std::vector<int> & gpus,
+    const std::vector<LayerSplitRange> & ranges,
+    const char * log_prefix);
+
+bool enable_layer_split_peer_access(
+    const std::vector<int> & gpus,
+    bool peer_access);
+
+bool init_layer_split_snapshot_backends(
+    const std::vector<LayerSplitShardMeta *> & shards,
+    std::vector<ggml_backend_t> & snapshot_backends,
+    const char * log_prefix);
+
+void free_layer_split_snapshot_backends(
+    const std::vector<LayerSplitShardMeta *> & shards,
+    std::vector<ggml_backend_t> & 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.

dflash/src/internal.h

@@ -22,6 +22,7 @@
 #include "ggml-backend.h"
 #include "gguf.h"
 
+#include "common/layer_split_utils.h"
 #include "dflash27b.h"
 
 namespace dflash::common {
@@ -176,12 +177,6 @@ inline bool is_eos_tok(int tok, const TargetWeights & w) {
         || (w.eos_id      >= 0 && tok == w.eos_id);
 }
 
-struct TargetLoadPlan {
-    int  layer_begin = 0;     // inclusive
-    int  layer_end   = -1;    // exclusive; <0 means all layers
-    bool load_output = true;  // output_norm + lm_head
-};
-
 // Load a Q4_K_M target model from a GGUF file on disk.
 // Returns false and sets last_error on failure.
 bool load_target_gguf(const std::string & path,
@@ -190,7 +185,7 @@ bool load_target_gguf(const std::string & path,
 
 bool load_target_gguf_partial(const std::string & path,
                               ggml_backend_t backend,
-                              const TargetLoadPlan & plan,
+                              const LayerSplitLoadPlan & plan,
                               TargetWeights & out);
 
 void free_target_weights(TargetWeights & w);

dflash/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 <cinttypes>
 #include <cstdint>
@@ -239,13 +240,13 @@ float get_f32_or(const gguf_context * g, const char * key, float fallback) {
 bool load_target_gguf(const std::string & path,
                       ggml_backend_t       backend,
                       TargetWeights &      out) {
-    TargetLoadPlan plan;
+    LayerSplitLoadPlan plan;
     return load_target_gguf_partial(path, backend, plan, out);
 }
 
 bool load_target_gguf_partial(const std::string & path,
                               ggml_backend_t       backend,
-                              const TargetLoadPlan & plan_in,
+                              const LayerSplitLoadPlan & plan_in,
                               TargetWeights &      out) {
 
     // ── 1. Parse metadata + create a ggml_context holding tensor descriptors ─
@@ -361,7 +362,7 @@ bool load_target_gguf_partial(const std::string & path,
         }
     }
 
-    TargetLoadPlan plan = plan_in;
+    LayerSplitLoadPlan plan = plan_in;
     if (plan.layer_begin < 0) plan.layer_begin = 0;
     if (plan.layer_end < 0) plan.layer_end = (int)n_layer;
     if (plan.layer_begin > plan.layer_end ||

dflash/src/qwen35/layer_split_daemon.cpp

@@ -14,8 +14,8 @@
 
 namespace dflash::common {
 
-bool run_target_layer_split_request(
-        std::vector<TargetLayerSplitShard> & shards,
+bool run_qwen35_layer_split_request(
+        std::vector<Qwen35LayerSplitShard> & 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<int32_t> 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)) {

dflash/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<TargetLayerSplitShard> & shards,
+bool run_qwen35_layer_split_request(
+        std::vector<Qwen35LayerSplitShard> & shards,
         DraftWeights * draft_weights,
         ggml_backend_t draft_backend,
         int draft_gpu,