Loading…
Token-count-based Batching: Faster, Cheaper Embedding Inference for Queries
MongoDBChengcheng Pei, Yuan Lin
Summary
Serving embedding models for short search queries often suffers from poor GPU efficiency because traffic is spiky and memory-bound. Traditional time-window and request-count batching strategies lead to inconsistent GPU utilization, while tensor padding wastes compute on empty tokens. To resolve these bottlenecks, Voyage AI implemented token-count-based batching paired with padding removal in inference engines like vLLM. The architecture uses Redis with Lua scripts to atomically aggregate pending requests until reaching an optimal hardware saturation token threshold. Across production deployments, this approach achieved a 50% reduction in GPU inference latency with 3X fewer GPUs and improved throughput by up to 8×.
Context
Embedding model inference for short queries in search and recommendation workloads is typically memory-bound, highly skewed in token length, and sensitive to latency (100–300 ms). Sequential serving and standard batching strategies based on time windows or request counts cause under- or over-filled batches and compute waste from tensor padding.
Approach / What changed
Voyage AI implemented token-count-based batching paired with padding removal in inference engines like vLLM. Requests are concatenated into variable-length super-sequences up to a hardware saturation threshold (such as 600 tokens for voyage-3 on A100). Model servers use Redis and Lua scripts to inspect queues and atomically claim batches bounded by total token count.
Takeaways
- Padding removal and variable-length processing allow inference engines to track actual token counts rather than padded batch dimensions, aligning GPU work with computation.
- Profiling inference latency versus token count reveals a saturation threshold (e.g., 600 tokens on A100 for voyage-3) where latency shifts from flat memory-bound overheads to linear scaling.
- Pairing token-count batching via Redis Lua scripts with vLLM achieved a 50% GPU inference latency reduction with 3X fewer GPUs and up to an 8x throughput improvement.
Related reading
MongoDB ·
MongoDB.local San Francisco 2026: Ship Production AI, Faster
MongoDB announced several features aimed at streamlining AI application development across its database and search ecosystem. The Voyage 4 embedding model family is now generally available, introducing cross-model compatibility within a shared embedding space alongside the video-capable voyage-multimodal-3.5 model. For in-database operations, MongoDB Community Edition introduced Automated Embedding in public preview to eliminate external synchronization pipelines for vector generation. Vector Search gained Lexical Prefilters, replacing the knnBeta operator with a dedicated vectorSearch operator and vector data type inside $search aggregations. Additionally, MongoDB made its intelligent assistant generally available in Compass and released the source code for its mongot search engine under the SSPL.
Peter RichardsMongoDB ·
Smarter AI Search, Powered by MongoDB Atlas and Pureinsights