lucebox/server/src/common/layer_split_runtime.cpp 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
#include "layer_split_runtime.h"

namespace luce::common {

bool run_layer_split_ar_decode(
        int last_tok,
        int committed,
        ThinkingBudget & budget,
        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) {
    const int n_gen = budget.remaining();
    if (n_gen <= 0) return true;

    std::vector<int32_t> history;
    if (sampler.needs_logit_processing()) {
        history.reserve(history_prefix.size() + out_tokens.size() + (size_t)n_gen);
        history.insert(history.end(), history_prefix.begin(), history_prefix.end());
        history.insert(history.end(), out_tokens.begin(), out_tokens.end());
        if ((int)prefill_last_logits.size() != vocab) return false;
        last_tok = sample_logits(prefill_last_logits.data(), vocab, sampler,
                                 history, rng);
    }

    if (io.is_cancelled()) {
        io.emit(-1);
        return true;
    }
    last_tok = budget.apply(last_tok).token;
    out_tokens.push_back(last_tok);
    if (sampler.needs_logit_processing()) history.push_back(last_tok);
    io.emit(last_tok);
    if (io.is_cancelled()) {
        io.emit(-1);
        return true;
    }
    if (is_eos(last_tok)) {
        io.emit(-1);
        return true;
    }
    ++committed;

    std::vector<float> logits_buf;
    for (int i = 1; i < n_gen; ++i) {
        std::vector<int32_t> one(1, last_tok);
        int next_tok = -1;
        logits_buf.clear();
        if (!forward_one(one, committed, next_tok,
                         sampler.needs_logit_processing() ? &logits_buf : nullptr)) {
            return false;
        }
        if (sampler.needs_logit_processing()) {
            if ((int)logits_buf.size() != vocab) return false;
            next_tok = sample_logits(logits_buf.data(), vocab, sampler,
                                     history, rng);
        }

        if (io.is_cancelled()) break;

        last_tok = budget.apply(next_tok).token;
        out_tokens.push_back(last_tok);
        if (sampler.needs_logit_processing()) history.push_back(last_tok);
        io.emit(last_tok);
        ++committed;
        if (io.is_cancelled()) break;
        if (is_eos(last_tok)) break;
    }

    io.emit(-1);
    return true;
}

}  // namespace luce::common