ts_wire/output.rs
1// Copyright (c) 2026 LightSeek Foundation
2//
3// Permission is hereby granted, free of charge, to any person obtaining a copy
4// of this software and associated documentation files (the "Software"), to deal
5// in the Software without restriction, including without limitation the rights
6// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7// copies of the Software, and to permit persons to whom the Software is
8// furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19// SOFTWARE.
20
21//! Per-step outputs from the engine to a frontend (`engine/io_struct.py`).
22
23use crate::{
24 codec::Dict,
25 error::{Error, Result},
26 value::Value,
27};
28
29wire_struct! {
30 /// One step's output for a frontend that detokenizes itself.
31 ///
32 /// Every column is parallel to `rids`. The load fields at the tail are
33 /// the scheduler's latest snapshot; an older producer omits them and
34 /// they decode as zero.
35 #[derive(Debug, Clone, PartialEq)]
36 pub struct BatchTokenIDOutSlim: array("BatchTokenIDOutSlim") {
37 /// Request ids. Always present on this wire.
38 pub rids: Option<Vec<String>> = None,
39 /// Token ids generated this step, per request.
40 pub output_ids: Vec<Vec<i32>>,
41 /// `""` while running, else the finish type (`"stop"`, `"length"`,
42 /// `"abort"`).
43 pub finished_reasons: Vec<String>,
44 /// Prompt length.
45 pub prompt_tokens: Vec<i64>,
46 /// Tokens generated so far.
47 pub completion_tokens: Vec<i64>,
48 /// Prompt tokens served from the prefix cache.
49 pub cached_tokens: Vec<i64>,
50 /// Logprobs of this step's tokens; empty for requests without logprobs.
51 pub output_token_logprobs_val: Vec<Vec<f64>>,
52 /// The token ids `output_token_logprobs_val` refers to.
53 pub output_token_logprobs_idx: Vec<Vec<i32>>,
54 /// The producing DP rank's engine index.
55 pub engine_index: i64 = 0,
56 /// Running requests.
57 pub num_running: i64 = 0,
58 /// Waiting requests.
59 pub num_waiting: i64 = 0,
60 /// KV pages in use by active requests.
61 pub kv_active_pages: i64 = 0,
62 /// KV pages in total.
63 pub kv_total_pages: i64 = 0,
64 /// The finish reason in full (`BaseFinishReason.to_json()`: the
65 /// matched token, an abort's message and codes), `None` while
66 /// running; one entry per request.
67 pub finish_reason_details: Vec<Option<FinishDetails>> = Vec::new(),
68 /// With speculative decoding, each request's verify rounds so far;
69 /// empty without.
70 pub spec_verify_ct: Vec<i64> = Vec::new(),
71 /// With speculative decoding, each finished request's accepted
72 /// tokens per verify round, `None` before it finishes; empty
73 /// without.
74 pub accept_draft_tokens: Vec<Option<f64>> = Vec::new(),
75 /// Prompt logprobs, shipped once per request: `None` for the first
76 /// position, then each position's logprob given its prefix. This and
77 /// the top-k columns hold one entry per request in a frame where
78 /// some request carries one, and are empty for the whole frame
79 /// otherwise.
80 pub input_token_logprobs_val: Vec<Vec<Option<f64>>> = Vec::new(),
81 /// The prompt tokens `input_token_logprobs_val` refers to.
82 pub input_token_logprobs_idx: Vec<Vec<i32>> = Vec::new(),
83 /// The top-k alternatives' logprobs at each prompt position (`None`
84 /// for the first), most likely first.
85 pub input_top_logprobs_val: Vec<Vec<Option<Vec<f64>>>> = Vec::new(),
86 /// Their token ids.
87 pub input_top_logprobs_idx: Vec<Vec<Option<Vec<i32>>>> = Vec::new(),
88 /// The top-k alternatives' logprobs of each token this step.
89 pub output_top_logprobs_val: Vec<Vec<Vec<f64>>> = Vec::new(),
90 /// Their token ids.
91 pub output_top_logprobs_idx: Vec<Vec<Vec<i32>>> = Vec::new(),
92 }
93}
94
95/// A finish reason in full, as a map (`"type"`, then `"matched"`,
96/// `"length"` or `"message"`/`"err_type"`/`"status_code"`).
97pub type FinishDetails = Dict<String, Value>;
98
99/// Top-k alternatives per token: their logprobs and token ids, most likely
100/// first.
101pub type TopLogprobs = (Vec<Vec<f64>>, Vec<Vec<i32>>);
102
103/// Top-k alternatives per prompt position, `None` for the first.
104pub type PromptTopLogprobs = (Vec<Option<Vec<f64>>>, Vec<Option<Vec<i32>>>);
105
106/// [`TopLogprobs`] borrowed from a batch.
107pub type TopLogprobRows<'a> = (&'a [Vec<f64>], &'a [Vec<i32>]);
108
109/// [`PromptTopLogprobs`] borrowed from a batch.
110pub type PromptTopLogprobRows<'a> = (&'a [Option<Vec<f64>>], &'a [Option<Vec<i32>>]);
111
112/// One request's slice of a [`BatchTokenIDOutSlim`].
113#[derive(Debug, Clone, Copy, PartialEq)]
114pub struct RequestOutput<'a> {
115 /// Request id.
116 pub rid: &'a str,
117 /// Token ids generated this step.
118 pub output_ids: &'a [i32],
119 /// The finish type, or `None` while the request runs.
120 pub finish_reason: Option<&'a str>,
121 /// Prompt length.
122 pub prompt_tokens: i64,
123 /// Tokens generated so far.
124 pub completion_tokens: i64,
125 /// Prompt tokens served from the prefix cache.
126 pub cached_tokens: i64,
127 /// Logprobs of this step's tokens; empty when not requested.
128 pub output_token_logprobs: &'a [f64],
129 /// The finish reason in full, once the request ended (from a producer
130 /// that sends it).
131 pub finish_details: Option<&'a FinishDetails>,
132 /// Verify rounds so far, with speculative decoding.
133 pub spec_verify_ct: Option<i64>,
134 /// Accepted tokens per verify round, once finished with speculative
135 /// decoding.
136 pub accept_draft_tokens: Option<f64>,
137 /// Prompt logprobs and their tokens, in the frame that ships them.
138 pub input_token_logprobs: Option<(&'a [Option<f64>], &'a [i32])>,
139 /// The prompt positions' top-k alternatives, in the frame that ships
140 /// them.
141 pub input_top_logprobs: Option<PromptTopLogprobRows<'a>>,
142 /// This step's tokens' top-k alternatives, when the request asked for
143 /// them.
144 pub output_top_logprobs: Option<TopLogprobRows<'a>>,
145}
146
147impl BatchTokenIDOutSlim {
148 /// The per-request rows, after checking that every column has one entry
149 /// per request.
150 pub fn requests(&self) -> Result<impl Iterator<Item = RequestOutput<'_>>> {
151 let rids = self.rids.as_deref().ok_or(Error::MissingField("rids"))?;
152 let len = rids.len();
153 let columns = [
154 ("output_ids", self.output_ids.len()),
155 ("finished_reasons", self.finished_reasons.len()),
156 ("prompt_tokens", self.prompt_tokens.len()),
157 ("completion_tokens", self.completion_tokens.len()),
158 ("cached_tokens", self.cached_tokens.len()),
159 (
160 "output_token_logprobs_val",
161 self.output_token_logprobs_val.len(),
162 ),
163 (
164 "output_token_logprobs_idx",
165 self.output_token_logprobs_idx.len(),
166 ),
167 ];
168 if let Some((name, found)) = columns.iter().find(|(_, found)| *found != len) {
169 return Err(Error::Invalid(format!(
170 "column `{name}` has {found} entries for {len} requests"
171 )));
172 }
173 // The tail's columns are either per request or empty for the frame.
174 let optional = [
175 ("finish_reason_details", self.finish_reason_details.len()),
176 ("spec_verify_ct", self.spec_verify_ct.len()),
177 ("accept_draft_tokens", self.accept_draft_tokens.len()),
178 (
179 "input_token_logprobs_val",
180 self.input_token_logprobs_val.len(),
181 ),
182 (
183 "input_token_logprobs_idx",
184 self.input_token_logprobs_idx.len(),
185 ),
186 ("input_top_logprobs_val", self.input_top_logprobs_val.len()),
187 ("input_top_logprobs_idx", self.input_top_logprobs_idx.len()),
188 (
189 "output_top_logprobs_val",
190 self.output_top_logprobs_val.len(),
191 ),
192 (
193 "output_top_logprobs_idx",
194 self.output_top_logprobs_idx.len(),
195 ),
196 ];
197 if let Some((name, found)) = optional
198 .iter()
199 .find(|(_, found)| *found != 0 && *found != len)
200 {
201 return Err(Error::Invalid(format!(
202 "column `{name}` has {found} entries for {len} requests"
203 )));
204 }
205 fn pair<'a, A, B>(
206 values: &'a [Vec<A>],
207 ids: &'a [Vec<B>],
208 i: usize,
209 ) -> Option<(&'a [A], &'a [B])> {
210 Some((values.get(i)?.as_slice(), ids.get(i)?.as_slice())).filter(|(v, _)| !v.is_empty())
211 }
212 Ok((0..len).map(move |i| RequestOutput {
213 rid: &rids[i],
214 output_ids: &self.output_ids[i],
215 finish_reason: Some(self.finished_reasons[i].as_str()).filter(|r| !r.is_empty()),
216 prompt_tokens: self.prompt_tokens[i],
217 completion_tokens: self.completion_tokens[i],
218 cached_tokens: self.cached_tokens[i],
219 output_token_logprobs: &self.output_token_logprobs_val[i],
220 finish_details: self.finish_reason_details.get(i).and_then(Option::as_ref),
221 spec_verify_ct: self.spec_verify_ct.get(i).copied(),
222 accept_draft_tokens: self.accept_draft_tokens.get(i).copied().flatten(),
223 input_token_logprobs: pair(
224 &self.input_token_logprobs_val,
225 &self.input_token_logprobs_idx,
226 i,
227 ),
228 input_top_logprobs: pair(
229 &self.input_top_logprobs_val,
230 &self.input_top_logprobs_idx,
231 i,
232 ),
233 output_top_logprobs: pair(
234 &self.output_top_logprobs_val,
235 &self.output_top_logprobs_idx,
236 i,
237 ),
238 }))
239 }
240}