nano-mooncake

English | 简体中文

A small, inspectable implementation of Mooncake-style KV-cache-centric LLM serving.

PD separation · paged KV cache · prefix reuse · shared-memory/TCP transfer · cache-aware scheduling · reproducible experiments

---

nano-mooncake is a self-contained systems project for learning and testing the core ideas behind disaggregated LLM serving. It makes the KV cache an explicit, transferable, persistent system object and provides a deterministic single-machine experiment platform for studying when cache reuse and prefill/decode separation help—or hurt.

The project favors clarity, measurable behavior, and reproducible evidence over production-scale throughput. The default model is a tiny randomly initialized HuggingFace decoder, so the generated text is not meaningful; the serving mechanics and measurements are the point.

Important

This repository is a functional single-machine research and learning platform, not a production inference server. See Scope and limitations before interpreting benchmark results.

Highlights

• Prefill/decode separation with both simple in-process and true multi-process execution paths.
• Paged KV cache backed by an explicit block pool and blockwise attention implementation.
• Global KV store with DRAM LRU, disk persistence, compatibility checks, atomic writes, and corrupt-entry quarantine.
• Prefix-aware scheduling based on cache locality, global cache availability, and worker queue pressure.
• Transfer abstraction with shared-memory and local TCP backends using the same snapshot codec.
• Serving semantics including bounded admission, backpressure, request timeouts, cooperative cancellation, streaming responses, and worker recovery policies.
• Credible experiment loop covering cold/warm cache behavior, worker-ratio pressure matrices, transfer overhead, queue wait, and fault injection.
• Evidence governance with pressure/fault budgets, decision reports, provenance manifests, reproducibility bundles, offline verification, and acceptance reports.
• Optional local HuggingFace tier for evaluating an existing pretrained checkpoint without downloading models at runtime.

Architecture

flowchart LR
    Client["Client / Benchmark"] --> API["API or Experiment Harness"]
    API --> Tokenizer["Tokenizer"]
    Tokenizer --> Scheduler["KV-aware Scheduler"]
    Scheduler --> Prefill["Prefill Worker"]
    Store["Mooncake Store<br/>DRAM + Disk"] <--> Prefill
    Prefill --> Snapshot["Paged KV Snapshot"]
    Snapshot --> Transfer["Shared Memory or TCP"]
    Transfer --> Decode["Decode Worker"]
    Decode --> Result["Tokens + Request Metrics"]
    Result --> Evidence["Pressure / Fault / Decision Gates"]
    Evidence --> Bundle["Verified Evidence Bundle"]

Loading

Request flow

1. The request is converted to token IDs.
2. The scheduler selects prefill and decode workers using prefix locality and queue pressure.
3. The prefill worker restores the longest reusable KV prefix, if available, and computes only the missing tail.
4. Full-prompt and block-boundary prefix snapshots are persisted in the store.
5. The resulting KV snapshot is transferred to a decode worker through shared memory or TCP.
6. The decode worker continues autoregressive generation from the transferred cache.
7. The runtime records TTFT, queue wait, cache hits, saved prefill tokens, transfer timing, and resource-lifecycle metrics.

Quick Start

Requirements

• Python 3.11+
• A local environment capable of installing PyTorch
• Linux is recommended for the multi-process and shared-memory experiment paths

Install

git clone git@github.com:BokaiGuo-Lincoln/nano-mooncake.git
cd nano-mooncake

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

python -m pip install --upgrade pip
python -m pip install -r requirements-dev.lock
python -m pip install -e .

Run the demo

nano-mooncake-demo

The module form works as well:

python -m nano_mooncake.demo

Run the tests

pytest -q

OpenAI-Compatible API

Start the local server:

nano-mooncake-api

Check health and queue metrics:

curl http://127.0.0.1:8000/healthz
curl http://127.0.0.1:8000/v1/metrics

Create a chat completion:

curl http://127.0.0.1:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "nano-mooncake",
    "messages": [
      {"role": "user", "content": "Explain KV-cache-aware scheduling."}
    ],
    "max_tokens": 8
  }'

Stream a response:

curl -N http://127.0.0.1:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "nano-mooncake",
    "messages": [
      {"role": "user", "content": "Explain prefill/decode separation."}
    ],
    "max_tokens": 8,
    "stream": true
  }'

The endpoint intentionally implements a small, practical subset of the OpenAI chat-completions schema rather than full protocol compatibility.

Experiments

Run the default deterministic experiment suite:

nano-mooncake-experiments \
  --output results/reports/single_machine_report.json

The suite compares the following execution modes:

Mode	PD split	Global KV reuse	Transfer path
Monolithic baseline	No	No	None
PD split, no global cache	Yes	No	Shared memory and TCP
PD split, global cache	Yes	Yes	Shared memory and TCP

It also exercises:

• high and low prefix reuse;
• exact prompt hits;
• long-prompt and short/long-decode scenarios;
• bounded-queue overload behavior;
• 1p1d, 2p1d, 1p2d, and 2p2d worker ratios;
• cold, warm, and DRAM-eviction-stressed cache lifecycles;
• worker crashes, timeouts, cancellation, shared-memory cleanup, and corrupt-cache recovery.

Metrics

Each request and aggregate summary can include:

• TTFT p50 / p95 / p99;
• total latency and decode latency;
• queue wait and worker utilization;
• cache source and cache-hit blocks;
• prefill tokens computed and saved;
• snapshot size and transfer bytes;
• encode, transport, decode, export, and import timing;
• leaked handles, residual queue slots, quarantined entries, and worker failures.

Results and Decision Views

Reports are registered in results/index.jsonl when experiments run locally. Render the accumulated history with:

# Comparable run summary
nano-mooncake-results

# Time-ordered changes
nano-mooncake-results --view trend

# Worker-ratio pressure matrix
nano-mooncake-results --view pressure

# Cache benefit after transfer accounting
nano-mooncake-results --view frontier

# Cache and PD-split win/loss conditions
nano-mooncake-results --view decision

# Fault and resource-lifecycle status
nano-mooncake-results --view faults

Views support JSON, Markdown-style terminal output, and CSV export:

nano-mooncake-results --view frontier --format json
nano-mooncake-results --view pressure --format csv \
  --output results/views/pressure.csv

Evidence and Acceptance Pipeline

The project distinguishes between a benchmark that ran and evidence that is suitable for comparison.

flowchart LR
    Run["Experiment Reports"] --> Registry["Result Registry"]
    Registry --> Pressure["Pressure Budget"]
    Registry --> Fault["Fault Budget"]
    Registry --> Decision["Frontier + Decision Report"]
    Pressure --> Acceptance["Acceptance Report"]
    Fault --> Acceptance
    Decision --> Acceptance
    Acceptance --> Export["Reproducibility Bundle"]
    Export --> Verify["Offline Digest Verification"]
    Verify --> Summary["Evidence Summary"]

Loading

Pressure and fault gates

nano-mooncake-pressure-budget \
  --index results/index.jsonl \
  --budget docs/pressure_budgets.json \
  --profile full \
  --output results/pressure_budget_full.json

nano-mooncake-fault-budget \
  --index results/index.jsonl \
  --output results/fault_budget_full.json

Isolated acceptance run

nano-mooncake-acceptance-run \
  --output-dir results/acceptance_runs/latest \
  --runs 2 \
  --max-new-tokens 2 \
  --pressure-profile full \
  --require-calibrated-pressure-history \
  --export-bundle

The acceptance runner writes an isolated index, experiment reports, acceptance verdict, diagnostics, artifact manifest, reproducibility bundle, ZIP archive, and verification reports.

Export and verify a bundle

nano-mooncake-export-bundle \
  --index results/index.jsonl \
  --output-dir results/repro_bundle \
  --pressure-profile full \
  --zip results/repro_bundle.zip

nano-mooncake-verify-bundle results/repro_bundle
nano-mooncake-verify-bundle results/repro_bundle.zip

Bundle verification recomputes file digests and canonical result-record digests, detecting missing, stale, modified, or inconsistent artifacts.

Optional Local HuggingFace Evidence

The default fast path uses a deterministic randomly initialized decoder. To evaluate an existing local pretrained checkpoint without network downloads:

nano-mooncake-local-hf-benchmark \
  --hf-model-path /path/to/local/hf/checkpoint \
  --output results/reports/local_hf_report.json

For a locally cached TinyLlama snapshot:

nano-mooncake-local-hf-benchmark \
  --preset tinyllama-cache \
  --output results/reports/local_hf_report.json \
  --skip-if-missing \
  --write-skipped

This tier performs native-decoder parity checks and records whether the evidence came from a real local checkpoint or the fast random-weight mechanism benchmark.

Command Reference

Command	Purpose
nano-mooncake-demo	Run a minimal local completion demo
nano-mooncake-api	Start the FastAPI server on 127.0.0.1:8000
nano-mooncake-experiments	Run the deterministic benchmark suite
nano-mooncake-results	Render compare, trend, pressure, frontier, decision, or fault views
nano-mooncake-rebuild-results	Rebuild the result registry from report files
nano-mooncake-pressure-budget	Evaluate or calibrate serving-pressure regression rules
nano-mooncake-fault-budget	Evaluate fault and cleanup invariants
nano-mooncake-acceptance-report	Produce a fixed-schema platform verdict
nano-mooncake-acceptance-run	Collect an isolated acceptance evidence set
nano-mooncake-acceptance-diagnostics	Rebuild diagnostics for an acceptance run
nano-mooncake-export-bundle	Export a reproducibility evidence bundle
nano-mooncake-verify-bundle	Verify a directory or ZIP bundle offline
nano-mooncake-evidence-summary	Produce a compact accepted/blocked evidence summary
nano-mooncake-decision-gap	Explain rows with insufficient decision evidence
nano-mooncake-pressure-diagnosis	Diagnose request-level pressure regressions
nano-mooncake-pressure-incident	Attribute failed pressure-budget incidents
nano-mooncake-pressure-margin	Explain near-budget pressure rules
nano-mooncake-local-hf-benchmark	Run the optional local-checkpoint evidence tier

Project Structure

nano-mooncake/
├── src/nano_mooncake/
│   ├── model.py                 # Decoder backends and paged attention
│   ├── kv_cache.py              # Block pool, paged cache, snapshots
│   ├── store.py                 # DRAM + disk KV cache store
│   ├── scheduler.py             # Cache- and queue-aware routing
│   ├── runtime.py               # Compact in-process request path
│   ├── runtime_modes.py         # Monolithic and multi-process runtimes
│   ├── transfer/                # Snapshot codec, shared memory, TCP
│   ├── api.py                   # Chat-completions API and admission queue
│   ├── experiments.py           # Workloads and experiment orchestration
│   ├── results_view.py          # Comparison and decision views
│   └── acceptance_*.py          # Acceptance and diagnostics pipeline
├── tests/                       # Unit, integration, serving, and fault tests
├── docs/
│   ├── workload_bundles.json    # Named workload definitions
│   ├── pressure_budgets.json    # Regression budget profiles
│   └── REMEDIATION_CHECKLIST.md # Completed hardening roadmap
├── .github/workflows/           # CI and scheduled evidence workflows
├── pyproject.toml
└── requirements-dev.lock

Design Principles

1. KV cache is a first-class object. It has an explicit layout, lifecycle, compatibility contract, transfer format, and persistence path.
2. Mechanisms should be inspectable. Python implementations are preferred where they make scheduling, transfer, and cache behavior easier to understand.
3. Comparisons require provenance. Reports include configuration, workload identity, source-tree digests, environment metadata, and evidence tier.
4. Hard gates and explanations are separate. Diagnostics explain failures or caveats but do not silently relax pressure or fault budgets.
5. Missing evidence is not a win. Decision views preserve insufficient_evidence rather than forcing every result into a positive or negative claim.

Scope and Limitations

This project currently provides:

• a single-machine Python serving prototype;
• real multi-process prefill/decode workers;
• real local shared-memory and TCP KV transfer paths;
• deterministic mechanism tests and experiment governance;
• optional evidence from an existing local HuggingFace checkpoint.

It does not currently provide:

• production-grade fused paged-attention kernels;
• GPU-optimized scheduling or memory management;
• RDMA or multi-node transfer;
• continuous batching comparable to production inference engines;
• distributed metadata and cache-coherence protocols;
• a claim of production throughput or latency parity with full Mooncake deployments.

Treat random-weight benchmark results as mechanism evidence, not model-quality or production-performance evidence.

Development

Install the locked development environment and editable package:

python -m pip install -r requirements-dev.lock
python -m pip install -e .

Run a focused test file during development:

pytest -q tests/test_runtime_modes.py

Run the complete suite before opening a pull request:

pytest -q

GitHub Actions also runs fault-injection checks, the complete test suite, a benchmark smoke, pressure/fault budget smokes, and bundle export verification.

Contributing

Contributions are welcome, especially when they preserve the project's inspectability and evidence discipline.

1. Fork the repository and create a focused branch.
2. Keep behavior changes small and include tests.
3. Avoid mixing refactors with feature changes.
4. Document new metrics, result fields, or claim boundaries.
5. Run pytest -q before submitting a pull request.

Useful contribution areas include transfer adapters, scheduler experiments, cache lifecycle tests, local-model evidence, visualization, and narrower production-inspired mechanisms that remain understandable on one machine.

Roadmap

The completed credibility-hardening roadmap is tracked in docs/REMEDIATION_CHECKLIST.md. Future work should build on the existing measurement and evidence surfaces rather than bypassing them.

Project Status

• Package version: 0.1.0
• Primary environment: Python 3.11+
• Current focus: single-machine serving mechanisms and credible experimental evidence
• Maturity: research/learning prototype

---

Built to make KV-cache-centric serving mechanics small enough to inspect and rigorous enough to measure.