|
| 1 | +// Generic server-facing backend for target layer split. |
| 2 | + |
| 3 | +#include "layer_split_backend.h" |
| 4 | + |
| 5 | +#include "io_utils.h" |
| 6 | + |
| 7 | +#include <chrono> |
| 8 | +#include <cstdio> |
| 9 | +#include <utility> |
| 10 | + |
| 11 | +namespace dflash::common { |
| 12 | + |
| 13 | +LayerSplitBackend::LayerSplitBackend(std::unique_ptr<LayerSplitAdapter> adapter) |
| 14 | + : adapter_(std::move(adapter)) {} |
| 15 | + |
| 16 | +LayerSplitBackend::~LayerSplitBackend() { shutdown(); } |
| 17 | + |
| 18 | +bool LayerSplitBackend::init() { |
| 19 | + if (!adapter_) { |
| 20 | + std::fprintf(stderr, "[target-split] missing model adapter\n"); |
| 21 | + return false; |
| 22 | + } |
| 23 | + return adapter_->init(); |
| 24 | +} |
| 25 | + |
| 26 | +void LayerSplitBackend::print_ready_banner() const { |
| 27 | + std::printf("[daemon] ready\n"); |
| 28 | + std::fflush(stdout); |
| 29 | +} |
| 30 | + |
| 31 | +bool LayerSplitBackend::park(const std::string & what) { |
| 32 | + std::fprintf(stderr, "[target-split] park is not supported yet (%s)\n", |
| 33 | + what.c_str()); |
| 34 | + return false; |
| 35 | +} |
| 36 | + |
| 37 | +bool LayerSplitBackend::unpark(const std::string & what) { |
| 38 | + std::fprintf(stderr, "[target-split] unpark is not supported yet (%s)\n", |
| 39 | + what.c_str()); |
| 40 | + return false; |
| 41 | +} |
| 42 | + |
| 43 | +GenerateResult LayerSplitBackend::run_from_state(const GenerateRequest & req, |
| 44 | + const DaemonIO & io, |
| 45 | + int base_pos, |
| 46 | + bool reset_state) { |
| 47 | + GenerateResult result; |
| 48 | + if (!adapter_) { |
| 49 | + result.error = "adapter"; |
| 50 | + return result; |
| 51 | + } |
| 52 | + |
| 53 | + DaemonIO out_io = io.with_token_callback(req.on_token); |
| 54 | + if (base_pos + (int)req.prompt.size() + req.n_gen + 1 > adapter_->max_context()) { |
| 55 | + result.error = "context"; |
| 56 | + return result; |
| 57 | + } |
| 58 | + if (req.do_sample && req.sampler.temp > 0.0f) { |
| 59 | + result.error = "sampling_unsupported"; |
| 60 | + return result; |
| 61 | + } |
| 62 | + |
| 63 | + adapter_->begin_request(req); |
| 64 | + if (reset_state) adapter_->reset_request_state(); |
| 65 | + |
| 66 | + const int prompt_len = (int)req.prompt.size(); |
| 67 | + int last_tok = (base_pos > 0 && prompt_len == 0) |
| 68 | + ? adapter_->current_last_token() |
| 69 | + : -1; |
| 70 | + int consumed = 0; |
| 71 | + auto t_prefill_start = std::chrono::steady_clock::now(); |
| 72 | + while (consumed < prompt_len) { |
| 73 | + int n_tokens = prompt_len - consumed; |
| 74 | + if (req.snap_pos >= 0 && req.snap_slot >= 0 && |
| 75 | + req.snap_pos > base_pos + consumed && |
| 76 | + req.snap_pos < base_pos + consumed + n_tokens) { |
| 77 | + n_tokens = req.snap_pos - (base_pos + consumed); |
| 78 | + } |
| 79 | + std::vector<int32_t> chunk(req.prompt.begin() + consumed, |
| 80 | + req.prompt.begin() + consumed + n_tokens); |
| 81 | + if (!adapter_->prefill(chunk, base_pos + consumed, last_tok)) { |
| 82 | + result.error = "prefill"; |
| 83 | + return result; |
| 84 | + } |
| 85 | + consumed += n_tokens; |
| 86 | + if (req.snap_pos >= 0 && req.snap_slot >= 0 && |
| 87 | + base_pos + consumed == req.snap_pos) { |
| 88 | + if (adapter_->snapshot_save(req.snap_slot)) { |
| 89 | + std::printf("[snap] inline slot=%d cur_pos=%d\n", |
| 90 | + req.snap_slot, req.snap_pos); |
| 91 | + std::fflush(stdout); |
| 92 | + } |
| 93 | + } |
| 94 | + } |
| 95 | + result.prefill_s = std::chrono::duration<double>( |
| 96 | + std::chrono::steady_clock::now() - t_prefill_start).count(); |
| 97 | + |
| 98 | + if (req.n_gen > 0) { |
| 99 | + if (last_tok < 0) { |
| 100 | + result.error = "decode_seed"; |
| 101 | + return result; |
| 102 | + } |
| 103 | + auto t_decode_start = std::chrono::steady_clock::now(); |
| 104 | + const bool ok = (base_pos == 0 && adapter_->can_dflash_decode()) |
| 105 | + ? adapter_->decode_dflash(req.prompt, base_pos, last_tok, req.n_gen, |
| 106 | + result.tokens, out_io) |
| 107 | + : adapter_->decode_ar(last_tok, base_pos + (int)req.prompt.size(), req.n_gen, |
| 108 | + result.tokens, out_io); |
| 109 | + if (!ok) { |
| 110 | + result.error = "decode"; |
| 111 | + return result; |
| 112 | + } |
| 113 | + result.decode_s = std::chrono::duration<double>( |
| 114 | + std::chrono::steady_clock::now() - t_decode_start).count(); |
| 115 | + } |
| 116 | + |
| 117 | + result.ok = true; |
| 118 | + return result; |
| 119 | +} |
| 120 | + |
| 121 | +GenerateResult LayerSplitBackend::generate(const GenerateRequest & req, |
| 122 | + const DaemonIO & io) { |
| 123 | + return run_from_state(req, io, /*base_pos=*/0, /*reset_state=*/true); |
| 124 | +} |
| 125 | + |
| 126 | +bool LayerSplitBackend::snapshot_save(int slot) { |
| 127 | + return adapter_ && adapter_->snapshot_save(slot); |
| 128 | +} |
| 129 | + |
| 130 | +void LayerSplitBackend::snapshot_free(int slot) { |
| 131 | + if (adapter_) adapter_->snapshot_free(slot); |
| 132 | +} |
| 133 | + |
| 134 | +bool LayerSplitBackend::snapshot_used(int slot) const { |
| 135 | + return adapter_ && adapter_->snapshot_used(slot); |
| 136 | +} |
| 137 | + |
| 138 | +int LayerSplitBackend::snapshot_cur_pos(int slot) const { |
| 139 | + return adapter_ ? adapter_->snapshot_cur_pos(slot) : 0; |
| 140 | +} |
| 141 | + |
| 142 | +GenerateResult LayerSplitBackend::restore_and_generate( |
| 143 | + int slot, const GenerateRequest & req, const DaemonIO & io) { |
| 144 | + GenerateResult result; |
| 145 | + if (!adapter_ || !adapter_->snapshot_restore(slot)) { |
| 146 | + result.error = "bad slot"; |
| 147 | + io.emit(-1); |
| 148 | + return result; |
| 149 | + } |
| 150 | + const int snap_pos = adapter_->snapshot_cur_pos(slot); |
| 151 | + if ((int)req.prompt.size() < snap_pos) { |
| 152 | + result.error = "snapshot_longer_than_prompt"; |
| 153 | + io.emit(-1); |
| 154 | + return result; |
| 155 | + } |
| 156 | + GenerateRequest delta_req = req; |
| 157 | + delta_req.prompt = std::vector<int32_t>( |
| 158 | + req.prompt.begin() + snap_pos, req.prompt.end()); |
| 159 | + return run_from_state(delta_req, io, snap_pos, /*reset_state=*/false); |
| 160 | +} |
| 161 | + |
| 162 | +ModelBackend::CompressResult |
| 163 | +LayerSplitBackend::compress(const CompressRequest & req) { |
| 164 | + return adapter_ ? adapter_->compress(req) : CompressResult{}; |
| 165 | +} |
| 166 | + |
| 167 | +bool LayerSplitBackend::handle_compress(const std::string & line, |
| 168 | + const DaemonIO & io) { |
| 169 | + std::string args = line.size() > 9 ? line.substr(9) : std::string{}; |
| 170 | + bool skip_park = false; |
| 171 | + const std::string suffix = " nopark"; |
| 172 | + if (args.size() >= suffix.size() && |
| 173 | + args.compare(args.size() - suffix.size(), suffix.size(), suffix) == 0) { |
| 174 | + skip_park = true; |
| 175 | + args.resize(args.size() - suffix.size()); |
| 176 | + } |
| 177 | + |
| 178 | + char ppath[1024]; |
| 179 | + int keep_x1000 = 0; |
| 180 | + char drafter_path[1024] = {0}; |
| 181 | + const int n = std::sscanf(args.c_str(), "%1023s %d %1023s", |
| 182 | + ppath, &keep_x1000, drafter_path); |
| 183 | + if (n < 2) { |
| 184 | + std::fprintf(stderr, "[target-split][compress] bad args\n"); |
| 185 | + io.emit(-1); |
| 186 | + return false; |
| 187 | + } |
| 188 | + |
| 189 | + CompressRequest req; |
| 190 | + req.input_ids = read_int32_file(ppath); |
| 191 | + req.keep_ratio = (float)keep_x1000 / 1000.0f; |
| 192 | + if (n >= 3 && drafter_path[0]) { |
| 193 | + req.drafter_path = drafter_path; |
| 194 | + } else if (adapter_) { |
| 195 | + req.drafter_path = adapter_->default_compress_drafter_path(); |
| 196 | + } |
| 197 | + req.skip_park = skip_park; |
| 198 | + |
| 199 | + CompressResult result = compress(req); |
| 200 | + for (int32_t t : result.compressed_ids) io.emit(t); |
| 201 | + io.emit(-1); |
| 202 | + return result.ok; |
| 203 | +} |
| 204 | + |
| 205 | +void LayerSplitBackend::free_drafter() { |
| 206 | + if (adapter_) adapter_->free_drafter(); |
| 207 | +} |
| 208 | + |
| 209 | +bool LayerSplitBackend::supports_dflash_spec_decode() const { |
| 210 | + return adapter_ && adapter_->supports_dflash_spec_decode(); |
| 211 | +} |
| 212 | + |
| 213 | +DFlashTarget * LayerSplitBackend::dflash_target() { |
| 214 | + return adapter_ ? adapter_->dflash_target() : nullptr; |
| 215 | +} |
| 216 | + |
| 217 | +bool LayerSplitBackend::supports_remote_draft() const { |
| 218 | + return adapter_ && adapter_->supports_remote_draft(); |
| 219 | +} |
| 220 | + |
| 221 | +void LayerSplitBackend::shutdown() { |
| 222 | + if (adapter_) adapter_->shutdown(); |
| 223 | +} |
| 224 | + |
| 225 | +} // namespace dflash::common |
0 commit comments