#506
Streaming Usage Information
predibase/lorax 머지됨 2024-06-11
RustStreamingToken AccountingFeature
Streaming chat completion 응답에 token usage 정보(prompt_tokens, completion_tokens, total_tokens)를 반환하도록 했다.
변경 사항
- Streaming chat completion 응답의 마지막 chunk에 usage 필드 추가
- OpenAI의 stream_options.include_usage 동작과 호환되는 토큰 사용량 리포팅
이 변경이 중요한 이유
Self-hosted LLM 서빙에서도 API 호출당 토큰 사용량을 추적해야 비용 관리와 모니터링이 가능하다. 기존에는 stream 모드에서 usage 정보가 누락되어 별도 토큰 카운팅 로직이 필요했다.
기술 상세
Rust router의 streaming response handler에서 generation 완료 시점에 prompt_tokens, completion_tokens, total_tokens를 집계하여 마지막 SSE chunk에 포함시켰다. OpenAI API의 stream_options 규격에 맞춘 구현이다. vLLM의 PagedAttention(Kwon et al., SOSP 2023)과 같은 서빙 시스템에서 토큰 사용량 추적은 KV 캐시 메모리 관리와 비용 모니터링의 기반이 된다.
관련 내부 작업
내부 LLM middleware에서 task/workspace 단위 토큰 사용량 추적에 활용 가능한 기반이 됐다. 다만 이 PR의 직접적인 내부 rollout artifact는 현재 확인되지 않는다.
관련 프로젝트 보기 →