Skip to content

Commit 3a817b4

Browse files
weicjweicj
committed
feat(server): add mixed-backend PFlash phase split
1 parent a40136e commit 3a817b4

19 files changed

Lines changed: 549 additions & 53 deletions

‎dflash/CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -240,6 +240,7 @@ add_library(dflash_common STATIC
240240
src/common/dflash_capture.cpp
241241
src/common/dflash_draft_ipc.cpp
242242
src/common/dflash_draft_ipc_daemon.cpp
243+
src/common/pflash_drafter_ipc.cpp
243244
src/common/dflash_draft_graph.cpp
244245
src/common/dflash_spec_decode.cpp
245246
src/qwen35/graph_builders.cpp
@@ -675,6 +676,7 @@ if(DFLASH27B_TESTS)
675676
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/src/ipc/backend_ipc_main.cpp")
676677
add_executable(backend_ipc_daemon
677678
src/ipc/backend_ipc_main.cpp
679+
src/common/pflash_drafter_ipc_daemon.cpp
678680
)
679681
target_include_directories(backend_ipc_daemon PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS})
680682
if(DFLASH27B_GPU_BACKEND STREQUAL "hip")

‎dflash/src/common/backend_factory.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,10 @@ bool arch_supports_remote_draft(const std::string & arch) {
2121
return arch == "qwen35";
2222
}
2323

24+
bool arch_supports_pflash_compression(const std::string & arch) {
25+
return arch == "qwen35" || arch == "qwen3";
26+
}
27+
2428
std::unique_ptr<ModelBackend> create_backend(const BackendArgs & args) {
2529
if (!args.model_path) {
2630
std::fprintf(stderr, "[backend_factory] model_path is null\n");

‎dflash/src/common/backend_factory.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,5 +65,6 @@ std::unique_ptr<ModelBackend> create_backend(const BackendArgs & args);
6565
std::string detect_arch(const char * model_path);
6666

6767
bool arch_supports_remote_draft(const std::string & arch);
68+
bool arch_supports_pflash_compression(const std::string & arch);
6869

6970
} // namespace dflash::common

‎dflash/src/common/backend_ipc.cpp‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ namespace dflash::common {
2222
const char * backend_ipc_mode_name(BackendIpcMode mode) {
2323
switch (mode) {
2424
case BackendIpcMode::DFlashDraft: return "dflash-draft";
25+
case BackendIpcMode::PFlashCompress: return "pflash-compress";
2526
}
2627
return "unknown";
2728
}
@@ -31,6 +32,10 @@ bool parse_backend_ipc_mode(const std::string & value, BackendIpcMode & out) {
3132
out = BackendIpcMode::DFlashDraft;
3233
return true;
3334
}
35+
if (value == "pflash-compress") {
36+
out = BackendIpcMode::PFlashCompress;
37+
return true;
38+
}
3439
return false;
3540
}
3641

‎dflash/src/common/backend_ipc.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ namespace dflash::common {
1818

1919
enum class BackendIpcMode {
2020
DFlashDraft,
21+
PFlashCompress,
2122
};
2223

2324
const char * backend_ipc_mode_name(BackendIpcMode mode);

‎dflash/src/common/model_backend.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,8 @@ struct ModelBackend {
143143
std::vector<int32_t> input_ids; // drafter-tokenized prompt
144144
float keep_ratio; // fraction to keep (0.0–1.0)
145145
std::string drafter_path; // GGUF path (for lazy-load)
146-
bool skip_park; // true on ≥32GB GPUs
146+
int drafter_gpu = 0; // backend-local GPU for PFlash drafter
147+
bool skip_park = false; // true on >=32GB GPUs
147148
};
148149

149150
struct CompressResult {

‎dflash/src/common/pflash_drafter_ipc.cpp‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
// pflash_drafter_ipc.cpp - PFlash drafter IPC client + daemon body.
2+
3+
#include "pflash_drafter_ipc.h"
4+
5+
#include <algorithm>
6+
#include <cmath>
7+
#include <cstdio>
8+
9+
namespace dflash::common {
10+
11+
bool PFlashDrafterIpcClient::start(
12+
const std::string & bin,
13+
const std::string & drafter_path,
14+
int drafter_gpu,
15+
const std::string & work_dir) {
16+
#if defined(_WIN32)
17+
(void)bin; (void)drafter_path; (void)drafter_gpu; (void)work_dir;
18+
std::fprintf(stderr, "PFlash drafter IPC is only implemented on POSIX hosts\n");
19+
return false;
20+
#else
21+
close();
22+
if (bin.empty() || drafter_path.empty()) return false;
23+
BackendIpcLaunchConfig launch;
24+
launch.bin = bin;
25+
launch.mode = BackendIpcMode::PFlashCompress;
26+
launch.payload_path = drafter_path;
27+
launch.work_dir = work_dir;
28+
launch.args.push_back("--draft-gpu=" + std::to_string(std::max(0, drafter_gpu)));
29+
if (!process_.start(launch)) {
30+
std::fprintf(stderr, "pflash-ipc backend process start failed\n");
31+
return false;
32+
}
33+
active_ = true;
34+
std::fprintf(stderr, "[pflash-ipc] ready bin=%s gpu=%d work_dir=%s\n",
35+
bin.c_str(), drafter_gpu, process_.work_dir().c_str());
36+
return true;
37+
#endif
38+
}
39+
40+
bool PFlashDrafterIpcClient::compress(
41+
const std::vector<int32_t> & input_ids,
42+
float keep_ratio,
43+
std::vector<int32_t> & compressed_ids) {
44+
#if defined(_WIN32)
45+
(void)input_ids; (void)keep_ratio; (void)compressed_ids;
46+
return false;
47+
#else
48+
compressed_ids.clear();
49+
FILE * cmd = process_.command_stream();
50+
const int stream_fd = process_.stream_fd();
51+
if (!active_ || !cmd || stream_fd < 0 || input_ids.empty()) return false;
52+
53+
const std::string path = process_.next_path("pflash_tokens");
54+
if (!write_int32_file(path, input_ids)) {
55+
std::fprintf(stderr, "pflash-ipc write tokens failed: %s\n", path.c_str());
56+
return false;
57+
}
58+
int keep_x1000 = (int)std::lround(std::max(0.0f, keep_ratio) * 1000.0f);
59+
keep_x1000 = std::max(0, std::min(1000, keep_x1000));
60+
61+
std::fprintf(cmd, "compress %d %s\n", keep_x1000, path.c_str());
62+
std::fflush(cmd);
63+
64+
int32_t status = -1;
65+
bool ok = read_exact_fd(stream_fd, &status, sizeof(status)) && status == 0;
66+
if (ok) {
67+
int32_t n_out = -1;
68+
ok = read_exact_fd(stream_fd, &n_out, sizeof(n_out)) && n_out > 0;
69+
if (ok) {
70+
compressed_ids.assign((size_t)n_out, 0);
71+
ok = read_exact_fd(stream_fd, compressed_ids.data(),
72+
compressed_ids.size() * sizeof(int32_t));
73+
}
74+
}
75+
std::remove(path.c_str());
76+
if (!ok) {
77+
std::fprintf(stderr, "pflash-ipc compress failed status=%d\n", status);
78+
compressed_ids.clear();
79+
close();
80+
}
81+
return ok;
82+
#endif
83+
}
84+
85+
void PFlashDrafterIpcClient::close() {
86+
process_.close();
87+
active_ = false;
88+
}
89+
90+
} // namespace dflash::common

‎dflash/src/common/pflash_drafter_ipc.h‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
// pflash_drafter_ipc.h - PFlash drafter IPC client + daemon entry.
2+
//
3+
// Used when target and PFlash drafter run on different compiled backends
4+
// (for example CUDA target + HIP drafter). The parent sends drafter-tokenized
5+
// prompt IDs to the daemon and receives compressed drafter token IDs back.
6+
7+
#pragma once
8+
9+
#include "backend_ipc.h"
10+
#include "io_utils.h"
11+
12+
#include <cstdint>
13+
#include <cstdio>
14+
#include <string>
15+
#include <vector>
16+
17+
namespace dflash::common {
18+
19+
class PFlashDrafterIpcClient {
20+
public:
21+
PFlashDrafterIpcClient() = default;
22+
PFlashDrafterIpcClient(const PFlashDrafterIpcClient &) = delete;
23+
PFlashDrafterIpcClient & operator=(const PFlashDrafterIpcClient &) = delete;
24+
~PFlashDrafterIpcClient() { close(); }
25+
26+
bool start(const std::string & bin,
27+
const std::string & drafter_path,
28+
int drafter_gpu,
29+
const std::string & work_dir);
30+
31+
bool compress(const std::vector<int32_t> & input_ids,
32+
float keep_ratio,
33+
std::vector<int32_t> & compressed_ids);
34+
35+
bool active() const { return active_; }
36+
void close();
37+
38+
private:
39+
BackendIpcProcess process_;
40+
bool active_ = false;
41+
};
42+
43+
int run_pflash_drafter_ipc_daemon(const char * drafter_path,
44+
int drafter_gpu,
45+
int stream_fd);
46+
47+
} // namespace dflash::common

‎dflash/src/common/pflash_drafter_ipc_daemon.cpp‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
// pflash_drafter_ipc_daemon.cpp - PFlash drafter IPC daemon body.
2+
3+
#include "pflash_drafter_ipc.h"
4+
5+
#include "dflash27b.h"
6+
#include "dflash_draft_ipc.h"
7+
#include "qwen3/qwen3_drafter.h"
8+
9+
#include <algorithm>
10+
#include <cstdio>
11+
#include <iostream>
12+
#include <sstream>
13+
14+
namespace dflash::common {
15+
16+
int run_pflash_drafter_ipc_daemon(const char * drafter_path,
17+
int drafter_gpu,
18+
int stream_fd) {
19+
#if defined(_WIN32)
20+
(void)drafter_path; (void)drafter_gpu; (void)stream_fd;
21+
std::fprintf(stderr, "PFlash drafter IPC daemon is only implemented on POSIX hosts\n");
22+
return 2;
23+
#else
24+
if (!drafter_path || stream_fd < 0) {
25+
std::fprintf(stderr,
26+
"usage: backend_ipc_daemon --backend-ipc-mode=pflash-compress <drafter.gguf> "
27+
"--stream-fd=FD [--draft-gpu=N]\n");
28+
return 2;
29+
}
30+
31+
DrafterContext ctx;
32+
if (!load_drafter(drafter_path, /*gpu_layers=*/999, std::max(0, drafter_gpu), ctx)) {
33+
std::fprintf(stderr, "[pflash-ipc-daemon] drafter load failed: %s\n",
34+
dflash27b_last_error());
35+
stream_status(stream_fd, -1);
36+
return 1;
37+
}
38+
39+
std::fprintf(stderr, "[pflash-ipc-daemon] ready gpu=%d\n", std::max(0, drafter_gpu));
40+
stream_status(stream_fd, 0);
41+
42+
std::string line;
43+
while (std::getline(std::cin, line)) {
44+
std::istringstream iss(line);
45+
std::string cmd;
46+
iss >> cmd;
47+
if (cmd == "quit" || cmd == "exit") break;
48+
if (cmd == "compress") {
49+
int keep_x1000 = 0;
50+
iss >> keep_x1000;
51+
std::string path = read_line_tail(iss);
52+
if (keep_x1000 < 0 || keep_x1000 > 1000 || path.empty()) {
53+
std::fprintf(stderr, "[pflash-ipc-daemon] bad compress: %s\n",
54+
line.c_str());
55+
stream_status(stream_fd, -1);
56+
continue;
57+
}
58+
auto input_ids = read_int32_file(path);
59+
if (input_ids.empty()) {
60+
std::fprintf(stderr, "[pflash-ipc-daemon] read tokens failed: %s\n",
61+
path.c_str());
62+
stream_status(stream_fd, -1);
63+
continue;
64+
}
65+
const float keep = (float)keep_x1000 / 1000.0f;
66+
auto compressed = drafter_score_and_compress(ctx, input_ids, keep);
67+
if (compressed.empty()) {
68+
std::fprintf(stderr, "[pflash-ipc-daemon] compress returned empty\n");
69+
stream_status(stream_fd, -1);
70+
continue;
71+
}
72+
const int32_t n_out = (int32_t)compressed.size();
73+
if (!stream_status(stream_fd, 0) ||
74+
!write_exact_fd(stream_fd, &n_out, sizeof(n_out)) ||
75+
!write_exact_fd(stream_fd, compressed.data(),
76+
compressed.size() * sizeof(int32_t))) {
77+
std::fprintf(stderr, "[pflash-ipc-daemon] stream write failed\n");
78+
break;
79+
}
80+
continue;
81+
}
82+
std::fprintf(stderr, "[pflash-ipc-daemon] unknown command: %s\n", line.c_str());
83+
stream_status(stream_fd, -1);
84+
}
85+
86+
free_drafter(ctx);
87+
std::fprintf(stderr, "[pflash-ipc-daemon] stopped\n");
88+
return 0;
89+
#endif
90+
}
91+
92+
} // namespace dflash::common

‎dflash/src/ipc/backend_ipc_main.cpp‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
#include "backend_ipc.h"
44
#include "dflash_draft_ipc.h"
5+
#include "pflash_drafter_ipc.h"
56

67
#include <algorithm>
78
#include <cstdio>
@@ -34,7 +35,10 @@ int main(int argc, char ** argv) {
3435
} else {
3536
std::fprintf(stderr,
3637
"usage: %s --backend-ipc-mode=dflash-draft <draft.safetensors|draft.gguf> "
37-
"--ring-cap=N --stream-fd=FD [--draft-gpu=N]\n",
38+
"--ring-cap=N --stream-fd=FD [--draft-gpu=N]\n"
39+
" or: %s --backend-ipc-mode=pflash-compress <drafter.gguf> "
40+
"--stream-fd=FD [--draft-gpu=N]\n",
41+
argv[0],
3842
argv[0]);
3943
return 2;
4044
}
@@ -64,6 +68,8 @@ int main(int argc, char ** argv) {
6468
switch (mode) {
6569
case BackendIpcMode::DFlashDraft:
6670
return run_dflash_draft_ipc_daemon(payload_path, ring_cap, draft_gpu, stream_fd);
71+
case BackendIpcMode::PFlashCompress:
72+
return run_pflash_drafter_ipc_daemon(payload_path, draft_gpu, stream_fd);
6773
}
6874
std::fprintf(stderr, "[backend-ipc-daemon] unsupported mode\n");
6975
return 2;

0 commit comments

Comments
 (0)