Hugging Face's logo

NYCU-MLLab
/
acdir-llada-math500

llada
discrete-diffusion
acdir
math
math500
critic-guided-decoding
Model card Files
xet
 Community
acdir-llada-math500
File size: 4,691 Bytes
4a28d4d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152


#include <algorithm>
#include <functional>
#include <memory>
#include <type_traits>
#include <utility>

#include "xgrammar/compiler.h"
#include "xgrammar/matcher.h"

#include "src/turbomind/engine/model_request.h"
#include "src/turbomind/engine/request.h"
#include "src/turbomind/utils/constant.h"
#include "src/turbomind/utils/metrics.h"

namespace turbomind {

ModelRequest::ModelRequest(Gateway* gateway, DataType data_type, int session_len, int vocab_size, int hidden_dim):
    gateway_{gateway},
    data_type_{data_type},
    session_len_{session_len},
    vocab_size_{vocab_size},
    hidden_dim_{hidden_dim}
{
}

void ModelRequest::Cancel()
{
    // request is finished if lock failed
    if (auto r = request_.lock()) {
        gateway_->cancel(std::move(r));
    }
}

void ModelRequest::End(std::function<void(int)> cb, uint64_t session_id)
{
    auto r = std::make_shared<Request>();

    r->id = r->session.id = session_id;
    r->session.kill_flag  = true;

    r->end_cb = std::move(cb);

    gateway_->kill(std::move(r));
}

auto ModelRequest::Forward(InputParam param, std::function<void()> cb) -> OutputParam
{
    inputs_  = std::make_shared<TensorMap>();
    outputs_ = std::make_shared<TensorMap>();

    auto add = [](auto& dest, auto key, auto dtype, auto where, auto shape, auto&&... dims) {
        Layout shape_;
        if constexpr (std::is_integral_v<decltype(shape)>) {
            shape_ = {shape, dims...};
        }
        else {
            shape_ = {shape.cbegin(), shape.cend()};
        }
        dest->emplace(key, Tensor{shape_, dtype, where});
    };

    auto& inputs = *param.tensors;

    TM_CHECK_EQ(inputs.at("input_ids").ndim(), 1);

    const int input_len  = inputs.at("input_ids").shape(0);
    const int output_len = param.gen_cfg.max_new_tokens;

    // Max possible length of a sequence, this depends on `history_len` which isn't available here, so `session_len`
    // is used instead
    const int max_seq_len = session_len_ + 1;
    const int max_out_len = std::min(output_len, session_len_) + 1;
    // This does not include histroy length in interactive mode
    const int max_in_out_len = std::min(input_len + output_len, session_len_) + 1;

    for (auto& [k, v] : *param.tensors) {
        inputs_->emplace(k, v);
    }

    add(outputs_, "output_ids", data_type_v<int>, kCPU, max_seq_len);
    add(outputs_, "sequence_length", data_type_v<int>, kCPU, 1);

    if (param.gen_cfg.output_logits) {
        const int len = param.gen_cfg.output_logits == GenerationConfig::kAll ? max_in_out_len : max_out_len;
        add(outputs_, "logits", data_type_, kCPU, len, vocab_size_);
    }

    if (param.gen_cfg.output_last_hidden_state) {
        const int len = param.gen_cfg.output_last_hidden_state == GenerationConfig::kAll ? max_in_out_len : max_out_len;
        add(outputs_, "last_hidden_state", data_type_, kCPU, len, hidden_dim_);
    }

    if (param.gen_cfg.output_logprobs) {
        add(outputs_, "logprob_vals", data_type_v<float>, kCPU, max_out_len, kMaxLogProb);
        add(outputs_, "logprob_indexes", data_type_v<int>, kCPU, max_out_len, kMaxLogProb);
        add(outputs_, "logprob_nums", data_type_v<int>, kCPU, max_out_len);
    }

    auto r = std::make_shared<Request>();

    for (const auto& [k, v] : *inputs_) {
        r->inputs.emplace(k, v);
    }
    for (const auto& [k, v] : *outputs_) {
        r->outputs.emplace(k, v);
    }

    auto state = std::make_shared<AtomicRequestState>();

    auto metrics = param.enable_metrics ? std::make_shared<RequestMetrics>() : nullptr;
    if (metrics) {
        metrics->enqueue_time.store(RequestMetrics::timestamp(), std::memory_order_relaxed);
        metrics->scheduled_time.store(0, std::memory_order_relaxed);
    }

    if (param.session.start_flag) {
        session_id_ = param.session.id;
    }

    r->id            = param.session.id;
    r->session       = param.session;
    r->gen_cfg       = param.gen_cfg;
    r->stream_output = param.stream_output;
    r->forward_cb    = std::move(cb);
    r->state         = state;
    r->metrics       = metrics;

    r->output_ids      = outputs_->at("output_ids");
    r->sequence_length = outputs_->at("sequence_length");

    if (grammar_) {
        r->grammar = std::move(grammar_);
        r->matcher = std::make_shared<xgrammar::GrammarMatcher>(*r->grammar);
    }

    // Keep a WEAK reference for canceling the request
    request_ = r;

    gateway_->push({std::move(r)});

    return OutputParam{outputs_, state, metrics};
}

void ModelRequest::setGrammar(const xgrammar::CompiledGrammar& grammar)
{
    grammar_ = std::make_shared<xgrammar::CompiledGrammar>(grammar);
}

}  // namespace turbomind