vLLM Continuous Batching & KV Cache GPU Concurrency

version: '3.8'

services:
  vllm-engine:
    image: vllm/vllm-openai:latest
    container_name: vllm_prod_engine
    restart: always
    runtime: nvidia
    ports:
      - "8000:8000"
    environment:
      - HUGGING_FACE_HUB_TOKEN=${HF_TOKEN}
      - NCCL_DEBUG=INFO
      - VLLM_LOGGING_LEVEL=INFO
    volumes:
      - /mnt/fast_nvme/hf_cache:/root/.cache/huggingface
      - /dev/shm:/dev/shm # Critical: Allocate host shared memory for PyTorch IPC
    shm_size: '32gb'
    command: >
      --model meta-llama/Llama-3.3-70B-Instruct
      --tensor-parallel-size 2
      --gpu-memory-utilization 0.95
      --max-model-len 8192
      --max-num-seqs 256
      --kv-cache-dtype fp8
      --port 8000
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 2
              capabilities: [gpu]
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
      interval: 10s
      timeout: 5s
      retries: 3

import asyncio
import aiohttp
import time

VLLM_URL = "http://localhost:8000/v1/chat/completions"
CONCURRENT_USERS = 64
MAX_TOKENS = 512

PROMPT = "Explain the difference between continuous batching and static batching in 3 paragraphs."

async def send_user_stream(session, user_id):
    payload = {
        "model": "llama-3.3-70b-instruct",
        "messages": [{"role": "user", "content": PROMPT}],
        "max_tokens": MAX_TOKENS,
        "temperature": 0.7,
        "stream": True,
    }
    
    start_time = time.time()
    ttft = None
    token_count = 0
    
    async with session.post(VLLM_URL, json=payload) as resp:
        async for chunk in resp.content:
            if not ttft:
                ttft = time.time() - start_time
            token_count += 1
            
    total_time = time.time() - start_time
    return {"user_id": user_id, "ttft": ttft, "total_time": total_time, "tokens": token_count}

async def main():
    print(f"Firing {CONCURRENT_USERS} simultaneous user streams into vLLM cluster...")
    async with aiohttp.ClientSession() as session:
        tasks = [send_user_stream(session, i) for i in range(CONCURRENT_USERS)]
        results = await asyncio.gather(*tasks)
        
    avg_ttft = sum(r["ttft"] for r in results) / len(results)
    total_tokens = sum(r["tokens"] for r in results)
    total_wall_time = max(r["total_time"] for r in results)
    
    print(f"\n--- Benchmark Results ---")
    print(f"Total Tokens Generated: {total_tokens}")
    print(f"Aggregate Throughput: {total_tokens / total_wall_time:.2f} tokens/sec")
    print(f"Average TTFT (Time-To-First-Token): {avg_ttft * 1000:.1f} ms")

if __name__ == "__main__":
    asyncio.run(main())