nanoSGLang

中文版

A minimal, high-performance LLM serving engine built from scratch in ~4000 lines of Python + PyTorch. Implements the core techniques from SGLang in a clean, educational codebase.

Tested on Qwen2.5-0.5B — reaches 3469 tok/s vs SGLang's 3201 tok/s on RTX A5000.

Architecture

nano_sglang/
  engine/
    engine.py           # Continuous batching loop, FlashInfer integration, CUDA Graph
    scheduler.py        # Batch scheduling with chunked prefill
    paged_kv_cache.py   # Block-based KV cache (FlashInfer NHD layout)
    radix_cache.py      # Prefix-sharing radix tree cache
    request.py          # Request lifecycle (WAITING -> RUNNING -> FINISHED)
    sampling.py         # Token sampling (FlashInfer fused / PyTorch fallback)
    kv_cache.py         # Naive contiguous KV cache (Phase 1)
    overlap.py          # CPU/GPU overlap scheduling (experimental)
    length_predictor.py # Output-length predictors for scheduling (experimental)
  model/
    causal_lm.py      # Llama/Qwen2 model implementation, fused RMSNorm
    tokenizer.py      # HuggingFace tokenizer wrapper
  server/
    api.py            # OpenAI-compatible REST API (FastAPI)
    metrics.py        # TTFT, TBT, throughput tracking
  decode/
    quantization.py   # FP8 / AWQ quantization
    speculative.py    # Speculative decoding (draft + verify)
    structured.py     # JSON-schema constrained generation
  distributed/
    tensor_parallel.py # Multi-GPU tensor parallelism
main.py               # Server entry point
benchmark.py           # Load testing tool
benchmarks/            # Correctness & throughput benchmarks
tests/                 # CPU-runnable unit tests

Quick Start

Install

pip install -r requirements.txt

# Required for best performance:
pip install flashinfer-python

# Optional (SM >= 80):
pip install flash-attn --no-build-isolation

Run Server

# Continuous batching mode (recommended):
python main.py --model-path Qwen/Qwen2.5-0.5B-Instruct \
    --num-blocks 8000 --max-batch-tokens 8192

# Naive mode (single request, no batching):
python main.py --model-path Qwen/Qwen2.5-0.5B-Instruct --naive

API Usage

# Chat completion
curl http://localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d '{
  "model": "qwen",
  "messages": [{"role": "user", "content": "Hello!"}],
  "max_tokens": 64,
  "stream": true
}'

# Text completion
curl http://localhost:8000/v1/completions -H "Content-Type: application/json" -d '{
  "model": "qwen",
  "prompt": "The capital of France is",
  "max_tokens": 32
}'

Offline Batch Inference

from benchmarks.nano_llm import NanoLLM, make_sampling_params

llm = NanoLLM(model_path="Qwen/Qwen2.5-0.5B-Instruct", num_blocks=8000)
results = llm.generate(
    prompt_token_ids=[[1, 2, 3], [4, 5, 6]],
    sampling_params=[make_sampling_params(max_tokens=64)] * 2,
)
for r in results:
    print(r.text)

Performance

Benchmark: 32 requests, prompt 100-512 tokens, output 64-128 tokens, Qwen2.5-0.5B, RTX A5000, greedy decoding.

System	Throughput (tok/s)
nanoSGLang (FlashInfer + fused ops)	3469
SGLang 0.5.6 (default, CUDA Graph ON)	3201
nanoSGLang (legacy, no FlashInfer)	177

Ablation

Measured by toggling each optimization individually on the same workload:

Configuration	tok/s	Delta
Legacy (flash_attn + per-token KV copy)	177	baseline
+ FlashInfer paged attention (zero-copy KV)	3076	+1642%
+ CUDA Graph (decode)	3077	+0%
+ Fused RMSNorm + Fused Sampling	3469	+13%

The dominant optimization is FlashInfer paged attention, which eliminates per-token, per-layer KV cache copy loops (O(layers × requests × seq_len) → O(1) via page table indirection). This accounts for ~97% of the total speedup.

CUDA Graph shows no measurable improvement on this 0.5B model — the model is too small for kernel launch overhead to matter. It would help on larger models (7B+).

Fused RMSNorm (via flashinfer.norm.fused_add_rmsnorm) and fused sampling (via flashinfer.sampling) contribute a real +13%.

Correctness

Greedy decoding (temperature=0) output comparison with SGLang on 6 prompts, Qwen2.5-0.5B (fused ops off):

#	Prompt	Status	nanoSGLang	SGLang
0	Hello, my name is	diverge @ token 1	John and I am a 20 year old male...	John**.** I am a student of the University...
1	The capital of France is	exact match	Paris. It is the largest city in Europe...	(same)
2	Write a Python function...	exact match	def sum_list(lst): return sum(lst)...	(same)
3	Question: What is 2 + 2?	exact match	4\nIs the above claim true?...	(same)
4	Once upon a time, in a small village,	diverge @ token 12	...Owl had a big family ...	...Owl had a special talent for predicting...
5	The three laws of robotics are: 1.	diverge @ token 19	...The law of the minimum time...	...The law of the minimum effort...

3/6 exact match. Divergences are due to bf16 numerical precision differences between attention backends — top-2 logits are very close at the divergence point, and small floating-point rounding differences flip the argmax. All outputs are semantically coherent. This is expected and consistent with SGLang's own behavior across different backends (e.g. FlashInfer vs Triton).

Features

• Continuous batching with mixed prefill + decode
• Chunked prefill for long prompts
• Paged KV cache with block-level memory management
• FlashInfer paged attention (zero-copy, NHD layout)
• CUDA Graph capture/replay for decode batches
• Fused operators (RMSNorm, sampling) via FlashInfer
• OpenAI-compatible API with streaming SSE
• Length-prediction scheduling (experimental): SJF admission + predicted-KV admission control

Length-Prediction Scheduling (experimental)

If the scheduler can guess how many tokens a request will generate, it can make better decisions than FCFS. Enabled with --enable-length-prediction (off by default — the FCFS path is untouched when disabled):

python main.py --model-path Qwen/Qwen2.5-0.5B-Instruct \
    --enable-length-prediction --length-predictor ewma

Design (no extra model needed):

• Pluggable predictors (--length-predictor): ewma (default) keeps an exponentially-weighted moving average of observed output lengths, bucketed by prompt-length bins; heuristic blends a running mean with a prompt-length-proportional term. Both clamp predictions by the request's max_tokens and update online as requests finish.
• Shortest-predicted-job-first admission: waiting requests with less predicted total work are admitted first, cutting average/median latency under queue pressure — the classic SJF result.
• Predicted-KV admission control: a request is only admitted if its predicted final KV footprint fits in the free block pool (with a no-deadlock fallback when nothing is running).
• Accuracy tracking: predicted-vs-actual MAE is reported under length_prediction in /metrics.

Measured on an A100 (SDPA legacy attention path), Qwen2.5-0.5B, 48 mixed-length requests (max_tokens cycling 8–128), client concurrency 24 against --max-running-requests 6 — i.e. a real waiting queue:

Metric	FCFS (off)	SJF (on, ewma)
TTFT p50	4.60 s	1.16 s (−75%)
E2E latency p50	5.47 s	3.79 s (−31%)
E2E latency p99	11.5 s	18.7 s (worse)
Throughput (tok/s)	102.6	108.3
Predictor MAE	—	~19 tokens

This is the textbook SJF trade-off: short requests stop waiting behind long ones (median and TTFT collapse), long requests pay at p99, per-token throughput is unchanged. Earlier runs on the slower naive-attention path showed the same shape (E2E p50 16.0 s → 5.8 s, avg −12%). When the server is not queue-limited (concurrency below --max-running-requests), the feature has little to act on and results match FCFS within noise.

Tested Models

• Qwen2.5-0.5B (verified correctness and throughput)

The model implementation follows the Llama/Qwen2 architecture. Other models sharing this architecture (e.g. Llama, Qwen2) may work but have not been tested.

Benchmarks

# Correctness (compare with SGLang, greedy decoding):
python benchmarks/bench_correctness.py --model-path Qwen/Qwen2.5-0.5B --system both

# Throughput:
python benchmarks/bench_throughput.py --model-path Qwen/Qwen2.5-0.5B --system nano

# Load testing:
python benchmark.py --url http://localhost:8000 --num-requests 100 --concurrency 32

# Mixed-length load (for length-prediction experiments):
python benchmark.py --url http://localhost:8000 --num-requests 48 --concurrency 24 --mixed-lengths

# CPU unit tests:
PYTHONPATH=. python3 tests/test_scheduler.py
PYTHONPATH=. python3 tests/test_length_prediction.py
PYTHONPATH=. python3 tests/test_speculative.py

Notable Fixes

Bugs found and fixed along the way (each is covered by a unit test in tests/):

• RoPE convention: apply_rope originally paired adjacent channels (interleaved); HF Llama/Qwen checkpoints use the half-split rotate_half pairing. Outputs were grammatical but degraded. Verified against HF transformers logits (tests/test_model_parity.py).
• KV-block exhaustion: running out of cache blocks crashed the forward pass and silently killed the engine loop, hanging every client. The scheduler now stalls decode steps under memory pressure and the engine aborts a failing batch instead of dying.
• Metrics counter leak: non-streaming requests never decremented num_waiting in /metrics.
• Speculative resampling: on rejection, resampling must draw from norm(max(0, q_target − p_draft)); the old formula reduced to resampling from the target itself, distorting the output distribution.
• FlashInfer import on Python < 3.10: flashinfer raises TypeError (dataclass slots), not ImportError; the fallback path now engages on any import failure.
• SDPA fallback on torch < 2.5: scaled_dot_product_attention only accepts enable_gqa= since torch 2.5; older versions now expand KV heads manually instead of crashing.
• Missing SamplingParams.ignore_eos: referenced by the engine and the offline benchmark wrapper but never defined.

Requirements

• Python >= 3.9 (FlashInfer itself needs >= 3.10; on older Pythons the engine automatically falls back to the legacy attention path)
• PyTorch >= 2.0
• CUDA GPU (SM >= 80 for FlashInfer / FlashAttention)
• flashinfer-python (for paged attention + fused ops)