Skip to content

Commit 6d2b9de

Browse files
weicjweicj
committed
feat(server): add native mixed-backend draft placement
1 parent 538bf53 commit 6d2b9de

12 files changed

Lines changed: 643 additions & 125 deletions

‎dflash/CMakeLists.txt‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -648,6 +648,27 @@ if(DFLASH27B_TESTS)
648648
endif()
649649
endif()
650650

651+
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/src/common/dflash_draft_ipc_main.cpp")
652+
add_executable(dflash_draft_ipc_daemon
653+
src/common/dflash_draft_ipc_main.cpp
654+
)
655+
target_include_directories(dflash_draft_ipc_daemon PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS})
656+
if(DFLASH27B_GPU_BACKEND STREQUAL "hip")
657+
target_compile_definitions(dflash_draft_ipc_daemon PRIVATE DFLASH27B_BACKEND_HIP=1 GGML_USE_HIP)
658+
else()
659+
target_compile_definitions(dflash_draft_ipc_daemon PRIVATE
660+
DFLASH27B_BACKEND_CUDA=1
661+
DFLASH27B_CUDA_MIN_SM=${_dflash_cuda_min_sm})
662+
endif()
663+
target_link_libraries(dflash_draft_ipc_daemon PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET} pthread)
664+
if(DFLASH27B_GPU_BACKEND STREQUAL "cuda")
665+
find_package(CUDAToolkit REQUIRED)
666+
target_link_libraries(dflash_draft_ipc_daemon PRIVATE CUDA::cudart)
667+
else()
668+
target_link_libraries(dflash_draft_ipc_daemon PRIVATE hip::host)
669+
endif()
670+
endif()
671+
651672
# Tokenizer test harness (no GPU needed — links static lib for tokenizer + GGUF reader)
652673
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_tokenizer_harness.cpp")
653674
add_executable(test_tokenizer_harness test/test_tokenizer_harness.cpp)

‎dflash/src/common/backend_factory.cpp‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@ std::string detect_arch(const char * model_path) {
1717
return info.arch;
1818
}
1919

20+
bool arch_supports_remote_draft(const std::string & arch) {
21+
return arch == "qwen35";
22+
}
23+
2024
std::unique_ptr<ModelBackend> create_backend(const BackendArgs & args) {
2125
if (!args.model_path) {
2226
std::fprintf(stderr, "[backend_factory] model_path is null\n");
@@ -37,7 +41,8 @@ std::unique_ptr<ModelBackend> create_backend(const BackendArgs & args) {
3741
cfg.target_path = args.model_path;
3842
cfg.draft_path = args.draft_path;
3943
cfg.device = args.device;
40-
cfg.draft_gpu = args.draft_gpu;
44+
cfg.draft_gpu = args.draft_device.gpu;
45+
cfg.remote_draft = args.remote_draft;
4146
cfg.stream_fd = args.stream_fd;
4247
cfg.fa_window = args.fa_window;
4348
cfg.kq_stride_pad = args.kq_stride_pad;

‎dflash/src/common/backend_factory.h‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@
1111
#pragma once
1212

1313
#include "model_backend.h"
14-
#include "device_placement.h"
14+
#include "placement/placement_config.h"
15+
#include "placement/remote_draft_config.h"
1516

1617
#include <memory>
1718
#include <string>
@@ -30,7 +31,8 @@ struct BackendArgs {
3031

3132
// Device placement
3233
DevicePlacement device;
33-
int draft_gpu = 0;
34+
DevicePlacement draft_device;
35+
RemoteDraftConfig remote_draft;
3436

3537
// I/O — only used when running under daemon_loop (legacy). The new
3638
// server passes -1 and uses on_token callbacks instead.
@@ -62,4 +64,6 @@ std::unique_ptr<ModelBackend> create_backend(const BackendArgs & args);
6264
// Useful for early dispatch (e.g. printing which backend will be used).
6365
std::string detect_arch(const char * model_path);
6466

67+
bool arch_supports_remote_draft(const std::string & arch);
68+
6569
} // namespace dflash::common

‎dflash/src/common/device_placement.h‎

Lines changed: 3 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,6 @@
1-
// Device placement configuration for model backends.
2-
//
3-
// Describes which GPU(s) to use for a model. Supports:
4-
// - Single-GPU: just gpu field
5-
// - Multi-GPU layer-split: layer_split_gpus + optional weights
6-
// - Peer access between GPUs
1+
// Compatibility include for older common/device_placement.h users.
2+
// New C++ placement code should include placement/placement_config.h.
73

84
#pragma once
95

10-
#include <vector>
11-
12-
namespace dflash::common {
13-
14-
struct DevicePlacement {
15-
int gpu = 0; // primary GPU (single-GPU mode)
16-
17-
// Multi-GPU layer-split. Empty = single GPU mode.
18-
std::vector<int> layer_split_gpus; // GPU IDs for each shard
19-
std::vector<double> layer_split_weights; // proportional layer distribution (optional)
20-
21-
bool peer_access = false; // enable CUDA peer access between GPUs
22-
int max_ctx = 8192; // max KV cache context length
23-
24-
bool is_layer_split() const { return layer_split_gpus.size() > 1; }
25-
26-
int primary_gpu() const {
27-
return layer_split_gpus.empty() ? gpu : layer_split_gpus[0];
28-
}
29-
};
30-
31-
} // namespace dflash::common
6+
#include "placement/placement_config.h"

‎dflash/src/common/dflash_draft_ipc_main.cpp‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
// Standalone DFlash draft IPC daemon entry point.
2+
3+
#include "dflash_draft_ipc.h"
4+
5+
#include <algorithm>
6+
#include <cstdio>
7+
#include <cstdlib>
8+
#include <cstring>
9+
10+
using namespace dflash::common;
11+
12+
int main(int argc, char ** argv) {
13+
if (argc < 3 || std::strcmp(argv[1], "--draft-ipc-daemon") != 0) {
14+
std::fprintf(stderr,
15+
"usage: %s --draft-ipc-daemon <draft.safetensors|draft.gguf> "
16+
"--ring-cap=N --stream-fd=FD [--draft-gpu=N]\n",
17+
argv[0]);
18+
return 2;
19+
}
20+
21+
const char * draft_path = argv[2];
22+
int ring_cap = 4096;
23+
int draft_gpu = 0;
24+
int stream_fd = -1;
25+
for (int i = 3; i < argc; i++) {
26+
if (std::strncmp(argv[i], "--ring-cap=", 11) == 0) {
27+
ring_cap = std::atoi(argv[i] + 11);
28+
} else if (std::strcmp(argv[i], "--ring-cap") == 0) {
29+
if (i + 1 < argc) ring_cap = std::atoi(argv[++i]);
30+
} else if (std::strncmp(argv[i], "--draft-gpu=", 12) == 0) {
31+
draft_gpu = std::max(0, std::atoi(argv[i] + 12));
32+
} else if (std::strcmp(argv[i], "--draft-gpu") == 0) {
33+
if (i + 1 < argc) draft_gpu = std::max(0, std::atoi(argv[++i]));
34+
} else if (std::strncmp(argv[i], "--stream-fd=", 12) == 0) {
35+
stream_fd = std::atoi(argv[i] + 12);
36+
} else if (std::strcmp(argv[i], "--stream-fd") == 0) {
37+
if (i + 1 < argc) stream_fd = std::atoi(argv[++i]);
38+
} else {
39+
std::fprintf(stderr, "[draft-ipc-daemon] unknown option: %s\n", argv[i]);
40+
return 2;
41+
}
42+
}
43+
44+
return run_dflash_draft_ipc_daemon(draft_path, ring_cap, draft_gpu, stream_fd);
45+
}

‎dflash/src/common/model_backend.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,12 @@ struct ModelBackend {
174174
// supports_dflash_spec_decode() returns true. Default returns nullptr.
175175
virtual class DFlashTarget * dflash_target() { return nullptr; }
176176

177+
// Return true when the backend can route draft execution through the
178+
// common remote-draft IPC transport. Model families that do not implement
179+
// the DFlash feature boundary keep the default false and are rejected by
180+
// the server before startup.
181+
virtual bool supports_remote_draft() const { return false; }
182+
177183
// ── Cleanup ──────────────────────────────────────────────────────
178184
// Release all resources (weights, cache, snapshots, drafter).
179185
// Called by run_daemon() before returning.

‎dflash/src/placement/placement_backend.h‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
// Backend placement identifiers shared by C++ server/runtime code.
2+
3+
#pragma once
4+
5+
#include <string>
6+
7+
namespace dflash::common {
8+
9+
enum class PlacementBackend {
10+
Auto,
11+
Cuda,
12+
Hip,
13+
};
14+
15+
inline const char * placement_backend_name(PlacementBackend backend) {
16+
switch (backend) {
17+
case PlacementBackend::Auto: return "auto";
18+
case PlacementBackend::Cuda: return "cuda";
19+
case PlacementBackend::Hip: return "hip";
20+
}
21+
return "auto";
22+
}
23+
24+
inline bool parse_placement_backend(const std::string & value,
25+
PlacementBackend & out) {
26+
if (value == "auto") {
27+
out = PlacementBackend::Auto;
28+
return true;
29+
}
30+
if (value == "cuda") {
31+
out = PlacementBackend::Cuda;
32+
return true;
33+
}
34+
if (value == "hip") {
35+
out = PlacementBackend::Hip;
36+
return true;
37+
}
38+
return false;
39+
}
40+
41+
inline PlacementBackend compiled_placement_backend() {
42+
#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP)
43+
return PlacementBackend::Hip;
44+
#else
45+
return PlacementBackend::Cuda;
46+
#endif
47+
}
48+
49+
inline bool placement_backend_supported(PlacementBackend backend) {
50+
return backend == PlacementBackend::Auto ||
51+
backend == compiled_placement_backend();
52+
}
53+
54+
} // namespace dflash::common

‎dflash/src/placement/placement_config.h‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
// Device placement configuration for model backends.
2+
//
3+
// Describes which backend device(s) to use for a model. Supports:
4+
// - Single-GPU: backend + gpu fields, exposed as cuda:0 / hip:0 / auto:0
5+
// - Multi-GPU layer-split: one backend + layer_split_gpus + optional weights
6+
// - Peer access between GPUs
7+
8+
#pragma once
9+
10+
#include "placement_backend.h"
11+
12+
#include <cstdlib>
13+
#include <string>
14+
#include <vector>
15+
16+
namespace dflash::common {
17+
18+
struct DevicePlacement {
19+
PlacementBackend backend = PlacementBackend::Auto;
20+
int gpu = 0; // primary GPU (single-GPU mode)
21+
22+
// Multi-GPU layer-split. Empty = single GPU mode.
23+
std::vector<int> layer_split_gpus; // GPU IDs for each shard
24+
std::vector<double> layer_split_weights; // proportional layer distribution (optional)
25+
26+
bool peer_access = false; // enable CUDA/HIP peer access between GPUs
27+
int max_ctx = 8192; // max KV cache context length
28+
29+
bool is_layer_split() const { return layer_split_gpus.size() > 1; }
30+
31+
int primary_gpu() const {
32+
return layer_split_gpus.empty() ? gpu : layer_split_gpus[0];
33+
}
34+
};
35+
36+
inline std::string placement_device_name(const DevicePlacement & device) {
37+
return std::string(placement_backend_name(device.backend)) + ":" +
38+
std::to_string(device.primary_gpu());
39+
}
40+
41+
inline bool parse_placement_device(const std::string & value,
42+
DevicePlacement & out) {
43+
const std::size_t sep = value.find(':');
44+
if (sep == std::string::npos || sep == 0 || sep + 1 >= value.size()) {
45+
return false;
46+
}
47+
48+
PlacementBackend backend = PlacementBackend::Auto;
49+
if (!parse_placement_backend(value.substr(0, sep), backend)) {
50+
return false;
51+
}
52+
53+
const std::string gpu_text = value.substr(sep + 1);
54+
char * end = nullptr;
55+
long gpu = std::strtol(gpu_text.c_str(), &end, 10);
56+
if (end == gpu_text.c_str() || *end != '\0' || gpu < 0) {
57+
return false;
58+
}
59+
60+
out.backend = backend;
61+
out.gpu = static_cast<int>(gpu);
62+
out.layer_split_gpus.clear();
63+
out.layer_split_weights.clear();
64+
return true;
65+
}
66+
67+
inline bool parse_placement_device_list(const std::string & value,
68+
DevicePlacement & out) {
69+
if (value.empty()) return false;
70+
71+
std::vector<int> gpus;
72+
PlacementBackend backend = PlacementBackend::Auto;
73+
bool have_backend = false;
74+
75+
std::size_t begin = 0;
76+
while (begin < value.size()) {
77+
const std::size_t end = value.find(',', begin);
78+
const std::string item = value.substr(
79+
begin,
80+
end == std::string::npos ? std::string::npos : end - begin);
81+
if (item.empty()) return false;
82+
83+
DevicePlacement parsed;
84+
if (!parse_placement_device(item, parsed)) return false;
85+
if (!have_backend) {
86+
backend = parsed.backend;
87+
have_backend = true;
88+
} else if (parsed.backend != backend) {
89+
return false;
90+
}
91+
gpus.push_back(parsed.gpu);
92+
93+
if (end == std::string::npos) break;
94+
begin = end + 1;
95+
}
96+
97+
if (gpus.empty()) return false;
98+
out.backend = backend;
99+
out.gpu = gpus[0];
100+
out.layer_split_gpus = gpus.size() > 1 ? gpus : std::vector<int>{};
101+
out.layer_split_weights.clear();
102+
return true;
103+
}
104+
105+
} // namespace dflash::common

‎dflash/src/placement/remote_draft_config.h‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
// Remote draft execution configuration for mixed-backend target/draft placement.
2+
3+
#pragma once
4+
5+
#include <string>
6+
7+
namespace dflash::common {
8+
9+
struct RemoteDraftConfig {
10+
std::string ipc_bin;
11+
std::string work_dir;
12+
int ring_cap = 0;
13+
14+
bool enabled() const { return !ipc_bin.empty(); }
15+
bool has_aux_options() const { return !work_dir.empty() || ring_cap > 0; }
16+
};
17+
18+
} // namespace dflash::common

0 commit comments

Comments
 (0)