Skip to content
← Back to Open Source
#506

Streaming Usage Information

predibase/lorax Merged 2024-06-11
View on GitHub
RustStreamingToken AccountingFeature

Added token usage information (prompt_tokens, completion_tokens, total_tokens) to streaming chat completion responses.

What Changed

  • Added usage field to the last chunk of streaming chat completion responses
  • Token usage reporting compatible with OpenAI's stream_options.include_usage behavior

Why It Matters

Even in self-hosted LLM serving, per-API-call token usage tracking is essential for cost management and monitoring. Previously, usage information was missing in stream mode, requiring separate token counting logic.

Technical Detail

In the Rust router's streaming response handler, aggregated prompt_tokens, completion_tokens, and total_tokens at generation completion and included them in the last SSE chunk. Implementation follows OpenAI API's stream_options specification. In serving systems like vLLM's PagedAttention (Kwon et al., SOSP 2023), token usage tracking is foundational for KV cache memory management and cost monitoring.

Related Internal Work

Provided the foundation for task/workspace-level token usage tracking in the internal LLM middleware. However, direct internal rollout artifacts from this specific PR are not currently confirmed.

View related project →