File size: 1,466 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 |
#pragma once
#include <memory>
#include "src/turbomind/core/core.h"
#include "src/turbomind/engine/gateway.h"
namespace xgrammar {
class CompiledGrammar;
}
namespace turbomind {
class ModelRequest {
public:
virtual ~ModelRequest() = default;
ModelRequest(Gateway* gateway, DataType data_type, int session_len, int vocab_size, int hidden_dim);
// Cancel running request
void Cancel();
// Reset the channel to uninitailized state, calls `notify` when done
void End(std::function<void(int)> cb, uint64_t session_id);
struct InputParam {
std::shared_ptr<TensorMap> tensors;
SessionParam session;
GenerationConfig gen_cfg;
bool stream_output;
bool enable_metrics;
};
struct OutputParam {
std::shared_ptr<TensorMap> tensors;
std::shared_ptr<AtomicRequestState> state;
std::shared_ptr<RequestMetrics> metrics;
};
OutputParam Forward(InputParam param, std::function<void()> cb);
void setGrammar(const xgrammar::CompiledGrammar& grammar);
protected:
Gateway* const gateway_;
const DataType data_type_;
const int session_len_;
const int hidden_dim_;
const int vocab_size_;
uint64_t session_id_;
std::weak_ptr<Request> request_;
std::shared_ptr<TensorMap> inputs_;
std::shared_ptr<TensorMap> outputs_;
std::shared_ptr<xgrammar::CompiledGrammar> grammar_;
};
} // namespace turbomind
|