suryanshchattree/npa-ppo-train
0
Neural PagedAttention
RL environment for training AI agents to manage GPU KV Cache in LLM inference servers.
Motivation
Modern LLM servers are bottlenecked by memory. Standard heuristics like LRU fail under bimodal traffic because they cannot predict future requests or distinguish VIP from Free users. This environment trains RL agents to do both.
Environment Description
A deterministic Python simulator of PagedAttention hardware.
- 1 Tick = 1 token generated across all active requests
- GPU: 10,000 blocks (160,000 tokens)
- CPU: 50,000 blocks (800,000 tokens) — swap target
- Traffic: 80% Chatters, 20% Power Users, 30% Returning Users
Observation Space
20-dimensional normalized float array [0.0, 1.0] (trend: [-1.0, 1.0]).
Action Space
Discrete(18), indices 0–17.
Tasks
Baseline Scores
Setup
docker build -t neural-paged-attention ./server
docker run -p 7860:7860 neural-paged-attentionRunning Inference
export HF_TOKEN=your_token
export MODEL_NAME=TinyLlama/TinyLlama-1.1B-Chat-v1.0
python inference.py