| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | use std::future::Future; |
| 5 | use std::pin::Pin; |
| 6 | use std::sync::Arc; |
| 7 | use std::task::{Context, Poll, ready}; |
| 8 | |
| 9 | use enum_as_inner::EnumAsInner; |
| 10 | use futures::stream::FusedStream; |
| 11 | use futures::{Stream, StreamExt as _, pin_mut}; |
| 12 | use serde::{Deserialize, Serialize}; |
| 13 | use vllm_engine_core_client::protocol::logprobs::Logprobs; |
| 14 | use vllm_engine_core_client::protocol::output::{ |
| 15 | EngineCoreFinishReason, RequestSpecDecodeMetrics, StopReason, |
| 16 | }; |
| 17 | use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; |
| 18 | use vllm_engine_core_client::protocol::tensor::WireNdArray; |
| 19 | use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream}; |
| 20 | |
| 21 | use crate::error::Result; |
| 22 | use crate::inflight::RequestGuard; |
| 23 | use crate::request_metrics::{RequestMetricsTracker, current_unix_timestamp_secs}; |
| 24 | |
| 25 | /// Token usage metadata for one request. |
| 26 | #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] |
| 27 | pub struct TokenUsage { |
| 28 | /// Number of prompt tokens sent to the engine. |
| 29 | pub prompt_token_count: usize, |
| 30 | /// Number of output tokens generated. |
| 31 | pub output_token_count: usize, |
| 32 | /// Number of prompt tokens served from cache. |
| 33 | pub cached_token_count: usize, |
| 34 | } |
| 35 | |
| 36 | /// Final raw token output plus terminal stream metadata. |
| 37 | #[derive(Debug, Clone, PartialEq)] |
| 38 | pub struct CollectedGenerateOutput { |
| 39 | pub request_id: String, |
| 40 | pub prompt_token_ids: Vec<u32>, |
| 41 | pub prompt_logprobs: Option<Logprobs>, |
| 42 | pub prompt_token_id_logprobs: Option<WireNdArray>, |
| 43 | pub token_ids: Vec<u32>, |
| 44 | pub logprobs: Option<Logprobs>, |
| 45 | pub finish_reason: FinishReason, |
| 46 | pub usage: TokenUsage, |
| 47 | /// Connector-specific KV transfer parameters for disaggregated serving. |
| 48 | pub kv_transfer_params: Option<serde_json::Value>, |
| 49 | /// Connector-specific encoder cache transfer parameters for disaggregated |
| 50 | /// serving. |
| 51 | pub ec_transfer_params: Option<serde_json::Value>, |
| 52 | /// Sampling support sets aligned one-to-one with generated token positions. |
| 53 | pub sampling_mask: Option<SamplingMask>, |
| 54 | /// Per-request speculative-decoding metrics from the terminal output. |
| 55 | pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>, |
| 56 | } |
| 57 | |
| 58 | /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for |
| 59 | /// one request. |
| 60 | #[derive(Debug, Clone, PartialEq)] |
| 61 | pub struct GeneratePromptInfo { |
| 62 | /// Original prompt token IDs for this request. |
| 63 | pub prompt_token_ids: Arc<[u32]>, |
| 64 | /// Prompt logprobs returned by engine-core for scored prompt positions, |
| 65 | /// when requested. |
| 66 | pub prompt_logprobs: Option<Logprobs>, |
| 67 | /// Log probabilities of `prompt_logprob_token_ids`, when requested. |
| 68 | pub prompt_token_id_logprobs: Option<WireNdArray>, |
| 69 | } |
| 70 | |
| 71 | /// The reason a request finished. |
| 72 | /// |
| 73 | /// This is a higher-level abstraction over engine-core's finish and stop |
| 74 | /// reasons. |
| 75 | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, EnumAsInner)] |
| 76 | pub enum FinishReason { |
| 77 | /// Generation stopped for a stop string, stop token, or EOS. |
| 78 | /// |
| 79 | /// The inner stop reason is present for explicit stop strings or stop |
| 80 | /// tokens, and absent for EOS-driven stops. |
| 81 | Stop(Option<StopReason>), |
| 82 | /// `max_tokens` or `max_model_len` was reached. |
| 83 | Length, |
| 84 | /// The request was aborted by the client. |
| 85 | Abort, |
| 86 | /// A retryable request-level internal error occurred. |
| 87 | Error, |
| 88 | /// A repetitive token pattern was detected. |
| 89 | Repetition(Option<StopReason>), |
| 90 | } |
| 91 | |
| 92 | impl FinishReason { |
| 93 | /// Construct a stop finish reason caused by EOS rather than an explicit |
| 94 | /// stop string/token. |
| 95 | pub fn stop_eos() -> Self { |
| 96 | Self::Stop(None) |
| 97 | } |
| 98 | |
| 99 | /// Returns a human-readable string for this finish reason, used for metrics |
| 100 | /// and reporting. |
| 101 | pub fn as_str(&self) -> &'static str { |
| 102 | match self { |
| 103 | Self::Stop(_) => "stop", |
| 104 | Self::Length => "length", |
| 105 | Self::Abort => "abort", |
| 106 | Self::Error => "error", |
| 107 | Self::Repetition(_) => "repetition", |
| 108 | } |
| 109 | } |
| 110 | |
| 111 | /// If this is a stop finish reason, returns the inner stop reason if it |
| 112 | /// exists. |
| 113 | pub fn as_stop_reason(&self) -> Option<&StopReason> { |
| 114 | match self { |
| 115 | Self::Stop(stop_reason) => stop_reason.as_ref(), |
| 116 | Self::Repetition(stop_reason) => stop_reason.as_ref(), |
| 117 | _ => None, |
| 118 | } |
| 119 | } |
| 120 | |
| 121 | /// If this is a stop finish reason, returns the inner stop reason if it |
| 122 | /// exists. |
| 123 | pub fn into_stop_reason(self) -> Option<StopReason> { |
| 124 | match self { |
| 125 | Self::Stop(stop_reason) => stop_reason, |
| 126 | Self::Repetition(stop_reason) => stop_reason, |
| 127 | _ => None, |
| 128 | } |
| 129 | } |
| 130 | } |
| 131 | |
| 132 | fn finish_reason_from_engine( |
| 133 | finish_reason: Option<EngineCoreFinishReason>, |
| 134 | stop_reason: Option<StopReason>, |
| 135 | ) -> Option<FinishReason> { |
| 136 | finish_reason.map(|reason| match reason { |
| 137 | EngineCoreFinishReason::Stop => FinishReason::Stop(stop_reason), |
| 138 | EngineCoreFinishReason::Length => FinishReason::Length, |
| 139 | EngineCoreFinishReason::Abort => FinishReason::Abort, |
| 140 | EngineCoreFinishReason::Error => FinishReason::Error, |
| 141 | EngineCoreFinishReason::Repetition => FinishReason::Repetition(stop_reason), |
| 142 | }) |
| 143 | } |
| 144 | |
| 145 | /// Token and logprob output item returned by [`GenerateOutputStream`]. |
| 146 | /// |
| 147 | /// Original Python output reference: |
| 148 | /// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/outputs.py#L85-L143> |
| 149 | #[derive(Debug, Clone, PartialEq)] |
| 150 | pub struct GenerateOutput { |
| 151 | /// Unique ID of the request that produced this output. |
| 152 | pub request_id: String, |
| 153 | /// One-time prompt metadata emitted only on the first output for this |
| 154 | /// request. |
| 155 | pub prompt_info: Option<GeneratePromptInfo>, |
| 156 | /// Newly produced token IDs for this step. |
| 157 | pub token_ids: Vec<u32>, |
| 158 | /// Sample logprobs for the generated positions in this step. |
| 159 | pub logprobs: Option<Logprobs>, |
| 160 | /// Terminal finish reason, when this is the final output for the request. |
| 161 | pub finish_reason: Option<FinishReason>, |
| 162 | /// Number of prompt tokens served from cache, when reported by prefill stats. |
| 163 | pub cached_token_count: usize, |
| 164 | /// Connector-specific KV transfer parameters for disaggregated serving. |
| 165 | pub kv_transfer_params: Option<serde_json::Value>, |
| 166 | /// Connector-specific encoder cache transfer parameters for disaggregated |
| 167 | /// serving. |
| 168 | pub ec_transfer_params: Option<serde_json::Value>, |
| 169 | /// Sampling support sets aligned one-to-one with `token_ids`. |
| 170 | pub sampling_mask: Option<SamplingMask>, |
| 171 | /// Per-request speculative-decoding metrics, present on terminal outputs. |
| 172 | pub spec_decode_metrics: Option<RequestSpecDecodeMetrics>, |
| 173 | } |
| 174 | |
| 175 | impl GenerateOutput { |
| 176 | /// Returns the prompt token IDs when this output carries |
| 177 | /// [`GeneratePromptInfo`]. |
| 178 | /// |
| 179 | /// Only the first output for a request can return `Some`; all later outputs |
| 180 | /// return `None`. |
| 181 | pub fn prompt_token_ids(&self) -> Option<&Arc<[u32]>> { |
| 182 | self.prompt_info.as_ref().map(|info| &info.prompt_token_ids) |
| 183 | } |
| 184 | |
| 185 | /// Returns the prompt logprobs when this output carries |
| 186 | /// [`GeneratePromptInfo`]. |
| 187 | /// |
| 188 | /// Only the first output for a request can return `Some`; all later outputs |
| 189 | /// return `None`. |
| 190 | pub fn prompt_logprobs(&self) -> Option<&Logprobs> { |
| 191 | self.prompt_info.as_ref().and_then(|info| info.prompt_logprobs.as_ref()) |
| 192 | } |
| 193 | |
| 194 | /// Returns whether this output is terminal for the request. |
| 195 | pub fn finished(&self) -> bool { |
| 196 | self.finish_reason.is_some() |
| 197 | } |
| 198 | } |
| 199 | |
| 200 | #[cfg(any(test, feature = "test-util"))] |
| 201 | impl GenerateOutput { |
| 202 | /// Build a [`GenerateOutput`] for tests. |
| 203 | pub fn for_test( |
| 204 | prompt_token_ids: Option<Arc<[u32]>>, |
| 205 | token_ids: Vec<u32>, |
| 206 | finish_reason: Option<FinishReason>, |
| 207 | ) -> Self { |
| 208 | Self { |
| 209 | request_id: String::new(), |
| 210 | prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo { |
| 211 | prompt_token_ids: ids, |
| 212 | prompt_logprobs: None, |
| 213 | prompt_token_id_logprobs: None, |
| 214 | }), |
| 215 | token_ids, |
| 216 | logprobs: None, |
| 217 | finish_reason, |
| 218 | cached_token_count: 0, |
| 219 | kv_transfer_params: None, |
| 220 | ec_transfer_params: None, |
| 221 | sampling_mask: None, |
| 222 | spec_decode_metrics: None, |
| 223 | } |
| 224 | } |
| 225 | } |
| 226 | |
| 227 | /// Stream of per-request generate outputs for one request. |
| 228 | /// |
| 229 | /// - A normal termination of the stream represents a clean completion of the |
| 230 | /// request, including a client-initiated abort, which yields a final output |
| 231 | /// with `finish_reason = Abort` before the stream ends. |
| 232 | /// - For errors or unexpected engine-side closes, the stream terminates with an error. |
| 233 | pub struct GenerateOutputStream { |
| 234 | pending_prompt_info: Option<GeneratePromptInfo>, |
| 235 | raw_stream: EngineCoreOutputStream, |
| 236 | request_metrics: RequestMetricsTracker, |
| 237 | /// Removes this request's external→internal tracking edge on drop. Held for |
| 238 | /// its `Drop` side effect only; never read directly. |
| 239 | _request_guard: RequestGuard, |
| 240 | } |
| 241 | |
| 242 | impl GenerateOutputStream { |
| 243 | /// Create a new generate output stream by adapting one raw engine-core |
| 244 | /// output stream. |
| 245 | pub(crate) fn new( |
| 246 | prompt_token_ids: Arc<[u32]>, |
| 247 | raw_stream: EngineCoreOutputStream, |
| 248 | request_metrics: RequestMetricsTracker, |
| 249 | request_guard: RequestGuard, |
| 250 | ) -> Self { |
| 251 | Self { |
| 252 | pending_prompt_info: Some(GeneratePromptInfo { |
| 253 | prompt_token_ids, |
| 254 | prompt_logprobs: None, |
| 255 | prompt_token_id_logprobs: None, |
| 256 | }), |
| 257 | raw_stream, |
| 258 | request_metrics, |
| 259 | _request_guard: request_guard, |
| 260 | } |
| 261 | } |
| 262 | |
| 263 | /// Return the internal engine request ID bound to this stream. |
| 264 | pub fn request_id(&self) -> &str { |
| 265 | self.raw_stream.request_id() |
| 266 | } |
| 267 | } |
| 268 | |
| 269 | impl Stream for GenerateOutputStream { |
| 270 | type Item = Result<GenerateOutput>; |
| 271 | |
| 272 | fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { |
| 273 | let raw = match ready!(Pin::new(&mut self.raw_stream).poll_next(cx)) { |
| 274 | Some(Ok(raw)) => raw, |
| 275 | Some(Err(error)) => return Poll::Ready(Some(Err(error.into()))), |
| 276 | None => return Poll::Ready(None), |
| 277 | }; |
| 278 | |
| 279 | let received_at = current_unix_timestamp_secs(); |
| 280 | self.request_metrics.observe_output(raw.timestamp, received_at, &raw.output); |
| 281 | |
| 282 | let raw = raw.output; |
| 283 | |
| 284 | // Populate the one-time prompt info on the first output. |
| 285 | if let Some(info) = &mut self.pending_prompt_info |
| 286 | && info.prompt_logprobs.is_none() |
| 287 | { |
| 288 | info.prompt_logprobs = |
| 289 | raw.new_prompt_logprobs_tensors.map(|value| value.into_direct().unwrap()); |
| 290 | } |
| 291 | if let Some(info) = &mut self.pending_prompt_info |
| 292 | && info.prompt_token_id_logprobs.is_none() |
| 293 | { |
| 294 | info.prompt_token_id_logprobs = raw.prompt_token_id_logprobs; |
| 295 | } |
| 296 | |
| 297 | let logprobs = raw.new_logprobs.map(|value| value.into_direct().unwrap()); |
| 298 | let sampling_mask = raw.new_sampling_mask.map(|value| value.into_direct().unwrap()); |
| 299 | if let Some(mask) = sampling_mask.as_ref() |
| 300 | && mask.rows.len() != raw.new_token_ids.len() |
| 301 | { |
| 302 | return Poll::Ready(Some(Err(crate::Error::SamplingMaskTokenCountMismatch { |
| 303 | request_id: raw.request_id, |
| 304 | token_count: raw.new_token_ids.len(), |
| 305 | row_count: mask.rows.len(), |
| 306 | }))); |
| 307 | } |
| 308 | let cached_token_count = raw |
| 309 | .prefill_stats |
| 310 | .as_ref() |
| 311 | .map(|stats| stats.num_cached_tokens as usize) |
| 312 | .unwrap_or(0); |
| 313 | |
| 314 | let finish_reason = finish_reason_from_engine(raw.finish_reason, raw.stop_reason); |
| 315 | if let Some(finish_reason) = finish_reason.as_ref() { |
| 316 | self.request_metrics.record_finished(received_at, finish_reason.clone()); |
| 317 | } |
| 318 | |
| 319 | let output = GenerateOutput { |
| 320 | request_id: raw.request_id, |
| 321 | prompt_info: self.pending_prompt_info.take(), |
| 322 | token_ids: raw.new_token_ids, |
| 323 | logprobs, |
| 324 | finish_reason, |
| 325 | cached_token_count, |
| 326 | kv_transfer_params: raw.kv_transfer_params, |
| 327 | ec_transfer_params: raw.ec_transfer_params, |
| 328 | sampling_mask, |
| 329 | spec_decode_metrics: raw.spec_decode_metrics, |
| 330 | }; |
| 331 | |
| 332 | Poll::Ready(Some(Ok(output))) |
| 333 | } |
| 334 | } |
| 335 | |
| 336 | impl FusedStream for GenerateOutputStream { |
| 337 | fn is_terminated(&self) -> bool { |
| 338 | self.raw_stream.is_terminated() |
| 339 | } |
| 340 | } |
| 341 | |
| 342 | impl Drop for GenerateOutputStream { |
| 343 | fn drop(&mut self) { |
| 344 | if self.raw_stream.is_terminated() { |
| 345 | // Already terminated cleanly, no need to record abort metrics. |
| 346 | return; |
| 347 | } |
| 348 | |
| 349 | // If the user or the upper layer drops a live generate stream, |
| 350 | // `EngineCoreOutputStream::Drop` will trigger an engine-side abort. Record the |
| 351 | // matching terminal request metrics here so frontend-driven aborts are still |
| 352 | // visible as `finished_reason=...` instead of disappearing from observability |
| 353 | // entirely. |
| 354 | let finish_reason = match AbortCause::current() { |
| 355 | AbortCause::DroppedStream => FinishReason::Abort, |
| 356 | AbortCause::StopStringMatched => FinishReason::Stop(None), |
| 357 | }; |
| 358 | |
| 359 | self.request_metrics |
| 360 | .record_finished(current_unix_timestamp_secs(), finish_reason); |
| 361 | } |
| 362 | } |
| 363 | |
| 364 | #[allow(clippy::manual_async_fn, reason = "specify `Send` bound")] |
| 365 | #[easy_ext::ext(GenerateOutputStreamExt)] |
| 366 | impl<T: Stream<Item = Result<GenerateOutput>> + Send> T { |
| 367 | /// Collect the raw generate stream to completion and return the final token |
| 368 | /// output. |
| 369 | pub fn collect_output(self) -> impl Future<Output = Result<CollectedGenerateOutput>> + Send { |
| 370 | async move { |
| 371 | let stream = self; |
| 372 | pin_mut!(stream); |
| 373 | let mut prompt_token_ids = None; |
| 374 | let mut prompt_logprobs = None; |
| 375 | let mut prompt_token_id_logprobs = None; |
| 376 | let mut cached_token_count = 0; |
| 377 | let mut collected: Option<CollectedGenerateOutput> = None; |
| 378 | |
| 379 | while let Some(output) = stream.next().await.transpose()? { |
| 380 | cached_token_count = cached_token_count.max(output.cached_token_count); |
| 381 | let sampling_mask = output.sampling_mask; |
| 382 | if let Some(info) = output.prompt_info { |
| 383 | if prompt_token_ids.is_none() { |
| 384 | prompt_token_ids = Some(info.prompt_token_ids.to_vec()); |
| 385 | } |
| 386 | if prompt_logprobs.is_none() { |
| 387 | prompt_logprobs = info.prompt_logprobs; |
| 388 | } |
| 389 | if prompt_token_id_logprobs.is_none() { |
| 390 | prompt_token_id_logprobs = info.prompt_token_id_logprobs; |
| 391 | } |
| 392 | } |
| 393 | |
| 394 | if let Some(existing) = collected.as_mut() { |
| 395 | existing.token_ids.extend(output.token_ids); |
| 396 | if let Some(step_logprobs) = output.logprobs { |
| 397 | if let Some(collected_logprobs) = existing.logprobs.as_mut() { |
| 398 | collected_logprobs.positions.extend(step_logprobs.positions); |
| 399 | } else { |
| 400 | existing.logprobs = Some(step_logprobs); |
| 401 | } |
| 402 | } |
| 403 | if let Some(mut mask) = sampling_mask { |
| 404 | existing.sampling_mask.get_or_insert_default().rows.append(&mut mask.rows); |
| 405 | } |
| 406 | } else { |
| 407 | collected = Some(CollectedGenerateOutput { |
| 408 | request_id: output.request_id, |
| 409 | prompt_token_ids: prompt_token_ids.take().unwrap_or_default(), |
| 410 | prompt_logprobs: prompt_logprobs.take(), |
| 411 | prompt_token_id_logprobs: prompt_token_id_logprobs.take(), |
| 412 | token_ids: output.token_ids, |
| 413 | logprobs: output.logprobs, |
| 414 | finish_reason: FinishReason::Error, |
| 415 | usage: TokenUsage { |
| 416 | prompt_token_count: prompt_token_ids.as_ref().map_or(0, Vec::len), |
| 417 | output_token_count: 0, |
| 418 | cached_token_count, |
| 419 | }, |
| 420 | kv_transfer_params: None, |
| 421 | ec_transfer_params: None, |
| 422 | sampling_mask, |
| 423 | spec_decode_metrics: None, |
| 424 | }); |
| 425 | } |
| 426 | |
| 427 | if let Some(finish_reason) = output.finish_reason { |
| 428 | let mut collected = collected.expect("terminal output must exist"); |
| 429 | collected.finish_reason = finish_reason; |
| 430 | collected.usage = TokenUsage { |
| 431 | prompt_token_count: collected.prompt_token_ids.len(), |
| 432 | output_token_count: collected.token_ids.len(), |
| 433 | cached_token_count, |
| 434 | }; |
| 435 | collected.kv_transfer_params = output.kv_transfer_params; |
| 436 | collected.ec_transfer_params = output.ec_transfer_params; |
| 437 | collected.spec_decode_metrics = output.spec_decode_metrics; |
| 438 | if let Some(mask) = collected.sampling_mask.as_ref() |
| 439 | && mask.rows.len() != collected.token_ids.len() |
| 440 | { |
| 441 | return Err(crate::Error::SamplingMaskTokenCountMismatch { |
| 442 | request_id: collected.request_id, |
| 443 | token_count: collected.token_ids.len(), |
| 444 | row_count: mask.rows.len(), |
| 445 | }); |
| 446 | } |
| 447 | return Ok(collected); |
| 448 | } |
| 449 | } |
| 450 | |
| 451 | unreachable!("generate stream should yield an error instead of closing early") |
| 452 | } |
| 453 | } |
| 454 | } |