suryanshchattree/neural-paged-attention-env

Neural PagedAttention

A reinforcement-learning environment for training agents to manage the GPU KV Cache in LLM inference servers.

Neural PagedAttention is a deterministic, high-fidelity simulator of vLLM-style PagedAttention memory hardware. It exposes the cache eviction / swapping / admission decisions as a Gymnasium-compatible RL environment, and ships with a suite of baseline and learned agents (LRU, Random, Q-Learning, DQN, LLM, PPO) so static and dynamic policies can be benchmarked head-to-head on identical workloads.

A full design write-up lives in `blog.md`.

[image]

Why this exists

Modern LLM servers are bottlenecked by GPU memory. Static heuristics like LRU/FIFO are simple and well-understood, but they cannot anticipate future requests, distinguish VIP from Free users, or react to bimodal traffic. Existing research into learned eviction policies is fragmented — every paper builds its own simulator and metrics.

This project provides a shared, reproducible environment where:

• —Static baselines (LRU, FIFO, size-weighted) and learned policies (DQN, PPO, LLM-as-policy) run on identical scenarios.
• —Traffic patterns include short chatters, power users, returning users, and viral spikes.
• —Episodes are scored on a composite of throughput, latency efficiency, and cache residency.

It is built as an early contribution to OpenEnv, a framework for standardizing RL environments in LLM infrastructure.

Features

• —Deterministic simulator of PagedAttention with GPU + CPU swap tiers.
• —20-D observation space capturing memory pressure, queue depth, idle-cache statistics, and SLA wait-time pressure.
• —18 discrete actions spanning eviction, swapping, admission, rejection, preemption, and garbage collection.
• —Three difficulty tiers — easy, medium, hard — with different traffic models and SLAs.
• —Six bundled agents plug-and-play behind a single API:
• —lru — Least Recently Used heuristic
• —random — Random baseline
• —qlearning — Tabular Q-learning
• —neural — Deep Q-Network (PyTorch)
• —llm — Qwen2.5-3B-Instruct as a zero-shot policy
• —ppo — Unsloth-fine-tuned Qwen with REINFORCE/PPO + LoRA via MLX
• —FastAPI HTTP server (/api/simulate, /api/agents, /api/settings) for headless batch runs and dashboards.
• —Next.js dashboard under dashboard/ for visualizing tick-level telemetry.
• —Hugging Face Spaces deploy out of the box (this repo is a Docker SDK Space).

Repository layout

.
├── app.py                  # FastAPI entrypoint (served on port 7860)
├── inference.py            # Standalone inference / smoke test
├── models.py               # Pydantic schemas for observations & actions
├── train_dqn.py            # DQN training loop
├── run_lru_test.py         # Quick LRU sanity check
├── run_qagent_test.py      # Quick tabular Q-learning sanity check
├── agents/                 # Agent implementations
│   ├── LRUAgent/
│   ├── RandomAgent/
│   ├── QLearningAgent/
│   ├── NeuralAgent/        # DQN
│   ├── LLMAgent/           # Qwen2.5 zero-shot
│   └── UnslothPPOAgent/    # PPO + LoRA via MLX
├── server/                 # Core environment
│   ├── environment.py      # KVCacheEnvironment (reset / step / state)
│   ├── env_components/     # Constants, traffic generator, scoring
│   └── session_manager.py  # Multi-session support
├── hf_train_space/         # Hugging Face Space for training
├── hf_ppo_space/           # Hugging Face Space for PPO inference
├── dashboard/              # Next.js telemetry dashboard
├── colab/                  # Colab training notebooks
├── assets/                 # Diagrams used in blog.md
├── blog.md                 # Project write-up
├── openapi.yaml            # OpenAPI spec for the HTTP API
├── Dockerfile              # Hugging Face Space build
├── requirements.txt
└── pyproject.toml

Environment specification

Observation space

20-dimensional float32 array, normalized to [0.0, 1.0] (the memory_pressure_trend field is in [-1.0, 1.0]).

Action space

Discrete(18):

Tasks

A tick corresponds to one token generated across all active requests. The default GPU has 10,000 blocks (≈160k tokens) and the CPU swap tier has 50,000 blocks (≈800k tokens). Traffic mix defaults to ~80% chatters, ~20% power users, with ~30% returning users.

Scoring

Each episode produces a composite score over:

• —Throughput — fraction of arriving requests successfully completed.
• —Latency efficiency — generation time vs. queue-wait time per request.
• —Cache residency — for returning users, whether their historical cache survived.

Throughput is the dominant term; cache residency acts as a tie-breaker.

Quickstart

Prerequisites

• —Python ≥ 3.10
• —(Optional) Docker, for the Hugging Face Space build
• —(Optional) HF_TOKEN if you want to pull gated models

Local install

git clone https://github.com/<your-org>/NeuralPagedAttention.git
cd NeuralPagedAttention

python -m venv .venv
source .venv/bin/activate

pip install -r requirements.txt

Run the API server

python app.py
# or
uvicorn app:app --host 0.0.0.0 --port 7860 --reload

The server is now available at http://localhost:7860 with interactive Swagger docs at http://localhost:7860/docs.

Run with Docker

docker build -t neural-paged-attention .
docker run -p 7860:7860 neural-paged-attention

Smoke-test an agent

python run_lru_test.py
python run_qagent_test.py

Train the DQN agent

python train_dqn.py

Run inference with the LLM agent

export HF_TOKEN=your_token
export MODEL_NAME=TinyLlama/TinyLlama-1.1B-Chat-v1.0
python inference.py

HTTP API

Full schema lives in `openapi.yaml`. The most useful endpoints:

GET /api/agents

List the agents available to the simulator.

GET /api/settings · POST /api/settings

Inspect or override environment constants (GPU_TOTAL_BLOCKS, CPU_TOTAL_BLOCKS, TOKENS_PER_BLOCK, per-task max_ticks, TRAFFIC_SEED).

POST /api/simulate

Run a full headless simulation and return tick- and session-level telemetry.

curl -X POST http://localhost:7860/api/simulate \
  -H "Content-Type: application/json" \
  -d '{"agent": "neural", "task": "hard"}'

Response:

{
  "status": "success",
  "agent": "neural",
  "task": "hard",
  "session_logs": [ ... ],
  "tick_logs":   [ ... ]
}

Legacy single-env endpoints (/reset, /step, /state) remain available but are deprecated; prefer the session-scoped variants exposed via server/session_manager.py.

Baseline scores

The table below reports the environment final score (a composite of throughput, latency efficiency, and cache residency, clamped to (0.001, 0.999)). Random / LRU / Q-Learning / DQN were measured locally with the default environment constants. LLM and PPO scores were fetched from the same backend the dashboard uses (https://suryanshchattree-neural-paged-attention-env.hf.space/api/simulate), where their tick budget is capped to keep request latency manageable on a CPU-only Space.

Cumulative-score traces from the live dashboard make the gap between difficulty tiers easy to read: every agent finishes easy in positive territory, while every agent is in free fall by the end of hard.

[image]

[image]

Raw per-run telemetry (reward totals, ticks run, arrival/completion counts, crash flags, LLM parse failures) is saved to `results/benchmarks.json` and the full LLM/PPO tick-level traces are in `results/remote_runs/`.

A few takeaways:

• —`easy` rewards a near-do-nothing policy more than people expect. Even Random scores 0.887 because traffic stays under capacity. The environment doesn't really start testing a policy until medium.
• —`medium` is the difficulty cliff. Random and DQN both trip the crash penalty within the first 20 ticks, collapsing to the 0.001 floor. LRU's 0.253 is the strongest passing score.
• —LLM and PPO get higher scores than Random / DQN on the harder tiers but on much shorter horizons. Both crash within ~25–30 ticks on medium/hard on the hosted backend. The composite score is biased upward by the small denominator (few requests arrived before crash). PPO on medium is the exception: it sustains 138 ticks and 12 completed requests, the best LLM-class number on that tier.
• —The classical baseline is harder to beat than people assume. LRU's 0.977 on easy and 0.253 on medium are the bars a learned policy actually has to clear. None of the learned agents shipped here clear them on medium end-to-end.

To regenerate locally:

python scripts/run_benchmarks.py                     # random / lru / qlearning / neural
python scripts/fetch_remote_scores.py                # llm / ppo via the live HF Space
# or, with GPU + cached Qwen2.5:
python scripts/run_benchmarks.py --include-llm --include-ppo

Or run an individual agent against the live API:

curl -s -X POST https://suryanshchattree-neural-paged-attention-env.hf.space/api/simulate \
  -H "Content-Type: application/json" \
  -d '{"agent": "ppo", "task": "medium"}' \
  | jq '.session_logs[].final_score'

Live dashboard

A web dashboard for running and inspecting these simulations is hosted at:

[https://neural-paged-attention-dashboard-7q.vercel.app/](https://neural-paged-attention-dashboard-7q.vercel.app/)

It is a Vite + React + Recharts client that talks to the public Hugging Face Space backend at https://suryanshchattree-neural-paged-attention-env.hf.space/api.

What the dashboard shows

• —Configuration panel — override the global environment constants (deterministic traffic seed, GPU and CPU total blocks, tokens per block, max ticks per difficulty), then trigger a batch run with the START BATCH SIMULATION button. Leaving fields blank uses the backend defaults.
• —Multi-agent comparison panel — MULTI-AGENT COMPARISON (CUMULATIVE SCORE) plots the running cumulative score of each agent on the same axis so you can read off relative performance over time, not just at episode end.
• —Agent focus panel — AGENT FOCUS: lets you drill into a single agent and shows live Tick, Score, Total Reward, current Task, Action, GPU / CPU Util, Memory Trend, and Yield Preempt indicator.
• —Queue panels — separate Free Requests and VIP Requests charts surface the queue dynamics that drive the SLA-miss penalties.
• —Live log — per-tick action / reward / tick stream (log-action, log-reward, log-tick) and a Live Score running total.
• —Status badges — COMPLETED and CRASHED flags per session, matching the crashed field in the API's session_logs.

A completed easy run vs. a fully-crashed hard run, side by side on the dashboard's traffic / token / macro-stats panels:

[image]

[image]

Insights you can read off the dashboard

1. 1.Crash localization. The cumulative-score line for an agent flattens the moment its crashed flag flips. Random and DQN go flat almost immediately on medium / hard, while LRU keeps drifting and only hits the wall mid-episode.
2. 2.Memory-pressure inflection points. The Memory Trend and GPU Utilization charts reveal that crashes follow a sharp, concentrated 5-tick rise in GPU utilization rather than slow drift — agents that don't pre-empt aggressively lose the window.
3. 3.Throughput vs. latency trade-off. On the LRU run, you can see a tight queue (VIP Requests near zero) but periodic stalls in cumulative score during traffic spikes — that's the policy paying for swap latency to keep VIPs admitted.
4. 4.Policy fingerprints. The action stream in the live log makes each agent's signature obvious: Random is uniform, LRU concentrates on Swap Oldest / Evict Oldest once GPU > 85%, the Q-learning agent over-issues Reject actions, and DQN's behavior changes character around the same memory-pressure thresholds the LRU heuristic uses.
5. 5.Reward shape ≠ score shape. Cumulative reward can plunge into very negative territory (large SLA penalties) while final score remains > 0.2 — useful when comparing crashed-but-completed runs against early-terminated ones.

Limitations

This is an early-stage simulator and intentionally simplified.

• —No hardware-level timing effects, memory bandwidth, or PCIe transfer cost modeling.
• —Traffic patterns are heuristic (compound sine + spikes), not replayed production traces.
• —Not yet validated against a live vLLM deployment; simulator and real-world performance may diverge.

See the "What it does not do" section of `blog.md` for a longer discussion.

Links

• —🌐 Live dashboard: https://neural-paged-attention-dashboard-7q.vercel.app/
• —🤗 Hugging Face Space (backend): https://suryanshchattree-neural-paged-attention-env.hf.space
• —📓 Blog / design write-up: `blog.md`
• —📜 OpenAPI spec: `openapi.yaml`