lucebox/server/src/common/layer_split_runtime.h at main · Luce-Org/lucebox · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
// Shared target layer-split runtime helpers.
//
// Model adapters still own their partial loaders, graph builders, caches, and
// snapshot payloads. This file keeps shared adapter runtime flow in one place
// so new adapters do not copy the same shell.

#pragma once

#include "gguf_inspect.h"
#include "layer_split_utils.h"
#include "model_backend.h"
#include "placement/placement_config.h"
#include "sampler.h"

#include <cstdio>
#include <functional>
#include <random>
#include <vector>

namespace dflash::common {

struct LayerSplitRuntimeInit {
    const char * target_path = nullptr;
    const DevicePlacement * device = nullptr;
    const char * log_prefix = "target-split";
};

template <typename Shard>
bool init_layer_split_runtime(const LayerSplitRuntimeInit & cfg,
                              std::vector<Shard> & shards,
                              std::vector<ggml_backend_t> & snapshot_backends) {
    const char * log_prefix = cfg.log_prefix ? cfg.log_prefix : "target-split";
    if (!cfg.target_path || !cfg.device ||
        cfg.device->layer_split_gpus.size() < 2) {
        std::fprintf(stderr, "[%s] invalid layer-split config\n", log_prefix);
        return false;
    }

    const auto info = inspect_gguf_model_info(cfg.target_path);
    const int n_layer = info.n_layer;
    if (n_layer <= 0) {
        std::fprintf(stderr, "[%s] failed to inspect target layer count\n",
                     log_prefix);
        return false;
    }

    const auto ranges = compute_layer_ranges(
        n_layer,
        (int)cfg.device->layer_split_gpus.size(),
        cfg.device->layer_split_weights);
    if (ranges.size() != cfg.device->layer_split_gpus.size()) {
        std::fprintf(stderr,
            "[%s] bad layer split for %zu GPUs and %d layers\n",
            log_prefix, cfg.device->layer_split_gpus.size(), n_layer);
        return false;
    }

    shards.resize(cfg.device->layer_split_gpus.size());
    auto shard_metas = layer_split_shard_metas(shards);
    if (!init_layer_split_shard_metas(
            shard_metas, cfg.device->layer_split_gpus, ranges,
            log_prefix)) {
        return false;
    }
    for (size_t i = 0; i < shard_metas.size(); ++i) {
        shard_metas[i]->placement_backend = cfg.device->layer_split_backend(i);
    }

    (void)enable_layer_split_peer_access(
        cfg.device->layer_split_gpus, cfg.device->peer_access);

    return init_layer_split_snapshot_backends(
        shard_metas, snapshot_backends, log_prefix);
}

using LayerSplitForwardStep = std::function<bool(
    const std::vector<int32_t> & tokens,
    int committed,
    int & next_tok,
    std::vector<float> * logits_out)>;

bool run_layer_split_ar_decode(
    int last_tok,
    int committed,
    int n_gen,
    int vocab,
    const std::vector<float> & prefill_last_logits,
    const SamplerCfg & sampler,
    std::mt19937_64 & rng,
    const std::vector<int32_t> & history_prefix,
    const LayerSplitForwardStep & forward_one,
    const std::function<bool(int)> & is_eos,
    std::vector<int32_t> & out_tokens,
    const DaemonIO & io);

}  // namespace dflash::common