KV cache problem

• During demoing we store key and value of past tokens in cache
  • New token need to attend to them
• Decoding is memory bandwidth bound. It reads
  • model weights
  • KV cache
  • just to get back one token
• How to get good throughput
  • Batch many sequence in each step
  • Helps with single read of weights
• One model weight is solved, KV cache is the bottleneck for throughput
  • The number of sequences whose caches fit in leftover GPU memory sets your batch size.
• Before vLLM
  • Each request got one continuous chunk size for maximum possible sequence length
  • Issues
    • memory reserved for token that were never generated
    • gaps between chunk too small to reuse
    • request with same system prompt their own copy

Paged Attention

• PagedAttention = KV cache in pages
• OS idea
  • doesn’t allocate large chunk of RAM
  • Gets pages on demand
  • Makes RAM look bigger
  • Maintains mapping table
• vLLM
  • splits KV cache on blocks
  • blocks are not contiguous
  • Each sequence gets a block table for mapping
• Standard attention kernel will not work
  • key values are no longer contiguous
• vLLM ships PagedAttentionKernel
  • follows block table to gather block
  • It’s a memory manager and a kernel
  • Some cost of indirection compensated by increased batch size
• More advantages
  • Sharing is easy, system prompt, beam search
  • We can preempt some sequence, move them to CPU etc

Continuous Batching

• Shortcoming of static batching
  • Some sequence finish early, they need to wait for large one to finish
• Continuous batching makes scheduling decision at every decode step
  • step = generating one token
• end token can leave the batch, KV blocks are feed, new sequence can join batch
  • Mad possible by pagedAttention
• Once concern
  • Long pre-fill is compute heavy and can stall decoding
  • Solution – chunked pre-fill
    • mixing dunked pre-fill with token budge
  • newer version stopped distinguishing them and simply decided how many tokens each request gets to process
• Knobs to tune
  • Max sequence per batch
  • Max tokens per step ( token budget )
• Trade off / metrics
  • Bigger batch raise throughput but each step takes longer.
    • Increase TPOT – time per output token
  • TTFT – time to first token
    • depends on suing delay and how fast refill is scheduled

Just a python process

• vLLM is user space python process
• At startup
  • it measure GPU weight memory needed
  • allocates one large tensor for rest
    • after that plus within this tensor
• OS and GPU drives just see big allocation that never changes
  • They don’t know about block tables
• Scheduling and block management runs on CPU
  • At each step
    • Engine packs blocks tables into tensors and sends the to GPU
  • Downside
    • For small model CPU overhead can become a bottleneck
      • Newer vLLM address this by running engines core in a separate process from the API server
• PagedAttentionKernel follow that map