A from-scratch Prefill/Decode disaggregation inference engine for LLMs

Disaggregated inference separates the two phases of LLM generation — the compute-intensive prefill (processing the prompt) and the memory-bandwidth-bound decode (generating tokens one at a time) — onto dedicated GPUs. This avoids the mutual interference between the two phases that limits throughput on collocated deployments.

nanoPD implements the full stack: a custom paged KV cache, chunked prefill, a custom CUDA paged attention kernel, multi-GPU KV transfer, an adaptive router driven by an analytical cost model, and a Poisson-arrival benchmark suite. All three serving strategies — Collocated, Disaggregated, and Adaptive — are implemented and benchmarked.

┌─────────────────────────────────────────────────────────────┐
│                      CentralScheduler                       │
│   (dispatch, KV transfer coordination, path accounting)     │
└────────────┬─────────────────┬───────────────┬─────────────┘
             │                 │               │
     ┌───────▼──────┐  ┌───────▼──────┐  ┌────▼──────────┐
     │  Collocated  │  │   Prefill    │  │    Decode     │
     │   Worker     │  │   Worker     │  │    Worker     │
     │  (GPU 0)     │  │ (GPU 1 / 3)  │  │   (GPU 2)    │
     └──────────────┘  └──────┬───────┘  └────▲──────────┘
                              │    KV Transfer │
                              └────────────────┘
             ↑
      ┌──────┴──────┐
      │   Router    │  ← analytical cost model decides path per request
      └─────────────┘

Each incoming request is routed by an analytical cost model that estimates the end-to-end latency of both strategies on the current hardware and picks the cheaper one. All parameters (prefill speed, decode speed, interference coefficient, inter-GPU bandwidth) are measured live on the actual device at startup.

• Paged KV Cache — block-granular memory management with copy-on-write for beam search / speculative decoding forks
• Chunked Prefill — long prompts are split into configurable chunks and interleaved with decode steps, keeping GPU utilisation high
• Custom CUDA Paged Attention Kernel — hand-written CUDA kernel for gather-scatter attention over non-contiguous KV blocks
• Async KV Transfer — prefill→decode KV cache migration over a dedicated CUDA stream via pinned memory relay or P2P direct (NVLink), with overlap against the compute stream
• Adaptive Router — per-request routing decision from a hardware-fitted analytical cost model; no oracle, no offline training
• Output Length Predictor — online Bayesian predictor for output length, used by the router to estimate decode cost before generation starts
• Multi-Worker CentralScheduler — concurrent Collocated and Disaggregated pipelines on separate threads, with dynamic batch management
• Poisson Arrival Benchmark — realistic open-loop load test with configurable arrival rate, workload distribution, warmup, and drain phases

1. Clone the repo

git clone https://github.com/your-username/nanoPD.git
cd nanoPD

2. Build the CUDA extension (compiles for the GPU on the current machine)

cd nanoPD/paged_attention
pip install -e . --no-build-isolation
cd ../..

Requires: Python ≥ 3.10, PyTorch ≥ 2.1, CUDA ≥ 11.8, and NVCC in PATH.
The extension uses -arch=native and auto-detects the installed GPU's compute capability.

3. Install Python dependencies

pip install transformers scipy numpy matplotlib

Runs 5 prompts through Engine.generate() on a single GPU. Suitable for RTX 4060/4070/4080 with Qwen2-1.5B.

python examples/demo_collocated.py
# or specify a local path:
python examples/demo_collocated.py --model /path/to/Qwen2-1.5B --gpu 0 --max-new-tokens 300

Loading Qwen/Qwen2-1.5B on cuda:0 ...
Model loaded in 7.1s

[1/5] Prompt: What is the capital of France?
  Response (7 tokens, 1.04s, 6.7 tok/s):
    The capital of France is Paris.
...

Runs the three-step demo on 8× GPUs — profile → fit cost model → 60 s Poisson adaptive inference. Results are written to output/output.txt.

python examples/demo_multiGPU.py --model /path/to/Qwen3-8B

# Skip re-profiling if output/data/profile_data.pt already exists
python examples/demo_multiGPU.py --model /path/to/Qwen3-8B --skip-profile

# Tune load
python examples/demo_multiGPU.py --model /path/to/Qwen3-8B \
    --skip-profile --arrival-rate 0.3 --workload long

Default GPU assignment:

Output files:

The router estimates end-to-end latency for both strategies using four hardware-measured parameters:

Key insight — The routing decision reduces to comparing two costs that are both linear in prompt length:

Extra cost of disaggregated : transfer_rate × L          (pay to move KV cache)
Extra cost of collocated     : γ × L × (load/batch_thresh)  (pay for prefill interference)

Disaggregated wins when:  γ / transfer_rate > batch_thresh / system_load

On RTX 4090: γ / transfer_rate ≈ 7.6 → Disaggregated wins from system_load ≥ 3
On H20: γ / transfer_rate ≈ 346 → Disaggregated wins at virtually any non-zero load

The full formula and per-hardware analysis is in docs/en/04-cost_model_en.md.

Tested on Qwen3-8B with two hardware configurations.

On H20, Disaggregated matches Collocated on short prompts (both 4.9 s) because 392 GB/s P2P bandwidth makes KV transfer nearly free. On the 4090, the 12.9 GB/s pinned-relay bandwidth adds a visible overhead — exactly as the cost model predicts.

RTX 4090 × 8:

H20:

• Adaptive saturates at ~240 tok/s on the 4090 and ~175 tok/s on H20 at moderate arrival rates
• Collocated is competitive at low load but p99 tail latency degrades quickly as concurrency grows
• Disaggregated (serial implementation) plateaus at ~25–30 tok/s regardless of device — the bottleneck is the lack of concurrent decode batching in the serial benchmark path

More plots and analysis: docs/en/07-benchmark_en.md

nanoPD/                            ← repo root
├── .gitignore
├── README.md
├── disaggregated_inference_engine.md   ← high-level design notes
├── examples/
│   ├── demo_collocated.py         ← single-GPU demo (Qwen2-1.5B)
│   └── demo_multiGPU.py           ← 8× GPU full pipeline demo
└── nanoPD/                        ← source package
    ├── block_manager/
    │   ├── block_manager.py       ← BlockSpaceManager, PhysicalBlock
    │   └── sequence.py            ← Sequence, SequenceGroup, SequenceStatus
    ├── engine/
    │   ├── engine.py              ← Engine (scheduler loop + chunked prefill)
    │   ├── model_runner.py        ← ModelRunner with paged_forward hook
    │   └── scheduler.py           ← Scheduler (prefill/decode batching)
    ├── paged_attention/
    │   └── csrc/                  ← CUDA kernels (paged attention + KV store)
    ├── workers/
    │   ├── collocated_worker.py
    │   ├── prefill_worker.py
    │   ├── decode_worker.py
    │   └── kv_transfer.py         ← async KV migration (pinned relay + P2P)
    ├── router/
    │   ├── central_scheduler.py   ← CentralScheduler (multi-worker dispatch)
    │   ├── router.py              ← Router (wraps cost model + predictor)
    │   └── output_lenth_predictor.py
    ├── cost_model/
    │   ├── profiler.py            ← device micro-benchmarks
    │   ├── analytical.py          ← curve fitting + routing decision
    │   ├── params.json            ← RTX 4090 fitted parameters
    │   └── params_h20.json        ← H20 fitted parameters
    ├── benchmark/
    │   ├── benchmark.py           ← static serial benchmark
    │   ├── benchmark_poisson.py   ← Poisson arrival benchmark
    │   ├── sweep.py               ← automated sweep across arrival rates
    │   └── plot_benchmark.py      ← result visualisation
    └── docs/
        ├── en/                    ← English documentation (7 modules)
        └── zh/                    ← Chinese documentation (7 modules)

Each module has a dedicated deep-dive doc covering design rationale, data structures, algorithms, and worked examples.