generate.proto 5.63 KB
Newer Older
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
1
2
syntax = "proto3";

Nicolas Patry's avatar
Nicolas Patry committed
3
package generate.v2;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
4

Olivier Dehaene's avatar
Olivier Dehaene committed
5
service TextGenerationService {
6
7
    /// Model Info
    rpc Info (InfoRequest) returns (InfoResponse) {}
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
8
    /// Service discovery
Olivier Dehaene's avatar
Olivier Dehaene committed
9
    rpc ServiceDiscovery (ServiceDiscoveryRequest) returns (ServiceDiscoveryResponse) {}
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
10
    /// Empties batch cache
Olivier Dehaene's avatar
Olivier Dehaene committed
11
    rpc ClearCache (ClearCacheRequest) returns (ClearCacheResponse);
12
13
    /// Remove requests from a cached batch
    rpc FilterBatch (FilterBatchRequest) returns (FilterBatchResponse);
14
15
    /// Warmup the model and compute max cache size
    rpc Warmup (WarmupRequest) returns (WarmupResponse);
16
17
18
19
    /// Prefill batch and decode first token
    rpc Prefill (PrefillRequest) returns (PrefillResponse);
    /// Decode token for a list of prefilled batches
    rpc Decode (DecodeRequest) returns (DecodeResponse);
20
21
    /// Health check
    rpc Health (HealthRequest) returns (HealthResponse);
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
22
23
}

24
25
26
message HealthRequest {}
message HealthResponse {}

27
28
29
30
31
32
33
/// Empty request
message InfoRequest {}

message InfoResponse {
    bool requires_padding = 1;
    string dtype = 2;
    string device_type = 3;
34
    optional uint32 window_size = 4;
Nicolas Patry's avatar
Nicolas Patry committed
35
    uint32 speculate = 5;
36
37
}

Olivier Dehaene's avatar
Olivier Dehaene committed
38
39
40
/// Empty request
message ServiceDiscoveryRequest {}

Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
41
message ServiceDiscoveryResponse {
Olivier Dehaene's avatar
Olivier Dehaene committed
42
    /// Other shards urls
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
43
44
45
    repeated string urls = 1;
}

46
47
48
49
message ClearCacheRequest {
    /// Optional batch id
    optional uint64 id = 1;
}
Olivier Dehaene's avatar
Olivier Dehaene committed
50
51
52
53

/// Empty response
message ClearCacheResponse {}

OlivierDehaene's avatar
OlivierDehaene committed
54
message NextTokenChooserParameters {
55
    /// exponential scaling output probability distribution
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
56
    float temperature = 1;
57
    /// restricting to the k highest probability elements
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
58
    uint32 top_k = 2;
59
    /// restricting to top tokens summing to prob_cut_off <= prob_cut_off
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
60
    float top_p = 3;
61
62
    /// restricting to top tokens summing to prob_cut_off <= prob_cut_off
    float typical_p = 4;
63
    /// apply sampling on the logits
64
    bool do_sample = 5;
65
    /// random seed for sampling
66
    uint64 seed = 6;
67
    /// repetition penalty
68
    float repetition_penalty = 7;
69
    /// token watermarking using "A Watermark for Large Language Models"
70
    bool watermark = 8;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
71
72
}

73
74
75
76
77
message StoppingCriteriaParameters {
    /// Maximum number of generated tokens
    uint32 max_new_tokens = 1;
    /// Optional stopping sequences
    repeated string stop_sequences = 2;
78
79
80
    /// Ignore end of sequence token
    /// used for benchmarking
    bool ignore_eos_token = 3;
81
82
}

Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
83
84
85
86
87
message Request {
    /// Request ID
    uint64 id = 1;
    /// The generation context
    string inputs = 2;
88
89
    /// Context truncation
    uint32 truncate = 3;
OlivierDehaene's avatar
OlivierDehaene committed
90
    /// Next Token Chooser Parameters
91
    NextTokenChooserParameters parameters = 4;
92
    /// Stopping Criteria Parameters
93
    StoppingCriteriaParameters stopping_parameters = 5;
94
95
    /// Return prefill logprobs
    bool prefill_logprobs = 6;
Nicolas Patry's avatar
Nicolas Patry committed
96
97
    /// Return most likely n tokens
    uint32 top_n_tokens = 7;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
98
99
100
101
102
103
104
}

message Batch {
    /// Batch ID
    uint64 id = 1;
    /// Individual requests
    repeated Request requests = 2;
Olivier Dehaene's avatar
Olivier Dehaene committed
105
106
    /// Batch size (==len(requests))
    uint32 size = 3;
107
108
    /// Maximum number of tokens this batch will grow to
    uint32 max_tokens = 4;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
109
110
}

111
112
113
114
115
116
117
118
119
120
121
message CachedBatch {
    /// Batch ID
    uint64 id = 1;
    /// Individual requests ids
    repeated uint64 request_ids = 2;
    /// Batch size (==len(requests))
    uint32 size = 3;
    /// Maximum number of tokens this batch will grow to
    uint32 max_tokens = 4;
}

122
123
124
125
126
127
enum FinishReason {
    FINISH_REASON_LENGTH = 0;
    FINISH_REASON_EOS_TOKEN = 1;
    FINISH_REASON_STOP_SEQUENCE = 2;
}

Olivier Dehaene's avatar
Olivier Dehaene committed
128
message GeneratedText {
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
129
    /// Output
130
    string text = 1;
131
    /// Number of generated tokens
132
    uint32 generated_tokens = 2;
133
    /// Finish reason
134
    FinishReason finish_reason = 3;
135
    /// Seed
136
    optional uint64 seed = 4;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
137
138
}

Nicolas Patry's avatar
Nicolas Patry committed
139
140
message Tokens {
    /// Token IDs
141
    repeated uint32 ids = 1;
Nicolas Patry's avatar
Nicolas Patry committed
142
    /// Logprobs
143
    repeated float logprobs = 2;
Nicolas Patry's avatar
Nicolas Patry committed
144
    /// tokens
145
    repeated string texts = 3;
Nicolas Patry's avatar
Nicolas Patry committed
146
147
    /// special
    repeated bool is_special = 4;
Nicolas Patry's avatar
Nicolas Patry committed
148
149
}

150
151
152
153
message Generation {
    /// Request ID
    uint64 request_id = 1;
    /// Prefill tokens (optional)
Nicolas Patry's avatar
Nicolas Patry committed
154
155
    Tokens prefill_tokens = 2;
    Tokens tokens = 3;
156
    /// Complete generated text
Nicolas Patry's avatar
Nicolas Patry committed
157
    optional GeneratedText generated_text = 4;
Nicolas Patry's avatar
Nicolas Patry committed
158
    /// Top tokens
Nicolas Patry's avatar
Nicolas Patry committed
159
    repeated Tokens top_tokens = 5;
160
161
}

162
163
164
165
message FilterBatchRequest {
    /// Batch ID
    uint64 batch_id = 1;
    /// Requests to keep
166
    repeated uint64 request_ids = 2;
167
168
169
170
}

message FilterBatchResponse {
    /// Filtered Batch (cached)
171
    CachedBatch batch = 1;
172
173
174
}


175
message PrefillRequest {
Olivier Dehaene's avatar
Olivier Dehaene committed
176
177
    /// Batch
    Batch batch = 1;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
178
179
}

180
181
182
message PrefillResponse {
    /// Generation
    repeated Generation generations = 1;
Olivier Dehaene's avatar
Olivier Dehaene committed
183
    /// Next batch (cached)
184
    optional CachedBatch batch = 2;
185
186
187
188
189
190
    /// Forward elapsed time in nanoseconds
    uint64 forward_ns = 3;
    /// Decode elapsed time in nanoseconds
    uint64 decode_ns = 4;
    /// Total elapsed time in nanoseconds
    uint64 total_ns = 5;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
191
192
}

193
message DecodeRequest {
Olivier Dehaene's avatar
Olivier Dehaene committed
194
    /// Cached batches
195
    repeated CachedBatch batches = 1;
Olivier Dehaene's avatar
Olivier Dehaene committed
196
}
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
197

198
199
200
message DecodeResponse {
    /// Decodes
    repeated Generation generations = 1;
Olivier Dehaene's avatar
Olivier Dehaene committed
201
    /// Next batch (cached)
202
    optional CachedBatch batch = 2;
203
204
205
206
207
208
209
210
    /// Forward elapsed time in nanoseconds
    uint64 forward_ns = 3;
    /// Decode elapsed time in nanoseconds
    uint64 decode_ns = 4;
    /// Total elapsed time in nanoseconds
    uint64 total_ns = 5;
    /// Concatenate elapsed time in nanoseconds
    optional uint64 concat_ns = 6;
211
}
212
213
214
215

message WarmupRequest {
    /// Batch to warmup on
    Batch batch = 1;
216
217
218
    uint32 max_input_length = 2;
    uint32 max_prefill_tokens = 3;
    uint32 max_total_tokens = 4;
219
220
221
}

/// Empty response
222
223
224
225
message WarmupResponse {
    /// Maximum number of tokens supported by the model
    optional uint32 max_supported_total_tokens = 1;
}