Skip to content

About

Nano vLLM

Resources

Stars

1 star

Watchers

0 watching

Forks

 
 

Latest commit

 

History

137 Commits

Folders and files

Repository files navigation

GeeeekExplorer%2Fnano-vllm | Trendshift

Nano-vLLM

A lightweight vLLM implementation built from scratch.

Key Features

  • 🚀 Fast offline inference - Comparable inference speeds to vLLM
  • 📖 Readable codebase - Clean implementation in ~ 1,200 lines of Python code
  • ⚡ Optimization Suite - Prefix caching, Tensor Parallelism, Torch compilation, CUDA graph, etc.

Sansho Fork

This fork is the serving target for Sansho models. See docs/sansho-port.md for the self-contained port plan: vendored Sansho runtime code, a SanshoModelRunner, YOCO-U native prefill/decode, and the path to continuous batching plus prefix caching.

The retained Sansho serving target is a local Hugging Face export such as outputs/eval/longrange-hf/swa256_control_500m. Run Sansho inference from the parent repo with this fork on PYTHONPATH, or run commands from this directory. For the retained SWA-256 export, the optimized Sansho runner is faster than local HF greedy inference at both c=1 and c=8 on GPU2.

Sansho Serving Roofline Anchors

The retained SWA-256 export is a roughly 351.6M parameter BF16 model (model.safetensors is 703,306,224 bytes). On the current RTX PRO 6000 Blackwell workstation, use 1.8 TB/s HBM bandwidth as a crude weight-streaming roofline. If each generated token streamed the weights exactly once and had no other overhead, the decode roof would be about 2,559 tok/s at c=1, 20.5k tok/s aggregate at c=8, and 163.8k tok/s aggregate at c=64.

Retained eager baseline, 512 input tokens and 128 generated tokens:

concurrency crude BW roof current decode share of roof peak reserved
c=1 2,559 tok/s 80.24 tok/s 3.1% 1,022 MiB
c=8 20,472 tok/s 601.62 tok/s 2.9% 1,528 MiB
c=64 163,776 tok/s 4,597.76 tok/s 2.8% 5,578 MiB

Opt-in runner CUDA graph replay is now implemented for full stable Sansho decode groups. With exact-shape graph warmup, sansho_decode_graph_replay_max_tokens=128, and the retained q=1 Triton self-attention decode kernel, the current 512/128 sweep is:

concurrency crude BW roof graph replay decode share of roof peak reserved
c=1 2,559 tok/s 242.64 tok/s 9.5% 888 MiB
c=8 20,475 tok/s 1,491.95 tok/s 7.3% 2,034 MiB
c=64 163,798 tok/s 6,171.71 tok/s 3.8% 11,234 MiB

This is a throughput win but not a memory-free one: the current implementation keeps one captured graph per decode position in the warmed segment. It uses a shared CUDA graph memory pool per replay chain, which keeps the retained c=64 shape under 12 GiB reserved instead of the earlier unpooled 70+ GiB. Use exact-shape warmup for eval or fixed-shape serving; leave the graph knob at 0 for dynamic serving shapes that cannot afford graph capture/pool residency.

This roofline is a sanity anchor, not a promise. It ignores activations, KV/state traffic, scheduler overhead, non-ideal HBM access patterns, sequential dependencies, and small-kernel launch overhead. It is still useful because the eager decode path is only around 3% of this crude bandwidth roof, while the current graph replay path reaches about 9.5% at c=1 and 7.3% at c=8.

The Sansho decode path also executes about 80 block calls per generated token: 60 self blocks (20 self layers times T=3), 16 cross-Mamba blocks, and 4 cross-attention blocks. An optimized dense serving stack of similar size would lean hard on CUDA graphs, fused kernels, and mature KV/decode plumbing, so those are first-class optimization targets here too. See PLAN.md for the item-by-item roofline punchlist and measurement contract.

Installation

pip install git+https://github.com/GeeeekExplorer/nano-vllm.git

Model Download

To download the model weights manually, use the following command:

huggingface-cli download --resume-download Qwen/Qwen3-0.6B \
  --local-dir ~/huggingface/Qwen3-0.6B/ \
  --local-dir-use-symlinks False

Quick Start

See example.py for usage. The API mirrors vLLM's interface with minor differences in the LLM.generate method:

from nanovllm import LLM, SamplingParams
llm = LLM("/YOUR/MODEL/PATH", enforce_eager=True, tensor_parallel_size=1)
sampling_params = SamplingParams(temperature=0.6, max_tokens=256)
prompts = ["Hello, Nano-vLLM."]
outputs = llm.generate(prompts, sampling_params)
outputs[0]["text"]

Sansho local export example:

cd /home/lhl/sansho/nano-vllm-sansho
CUDA_VISIBLE_DEVICES=2 mamba run -n sansho python - <<'PY'
from nanovllm import LLM, SamplingParams

llm = LLM(
    "../outputs/eval/longrange-hf/swa256_control_500m",
    max_model_len=8192,
    max_num_seqs=8,
    max_num_batched_tokens=16384,
    enforce_eager=True,
)
outputs = llm.generate(["Hello, Sansho."], SamplingParams(temperature=0.0, max_tokens=32), use_tqdm=False)
print(outputs[0]["text"])
llm.exit()
PY

Request selected-token logprobs from the Python API with SamplingParams.logprobs. The integer value controls how many top-token logprobs are retained per generated token; 0 keeps only the selected token logprob:

outputs = llm.generate(
    ["Hello, Sansho."],
    SamplingParams(temperature=0.0, max_tokens=8, logprobs=2),
    use_tqdm=False,
)
outputs[0]["logprobs"]
outputs[0]["top_logprobs"]

scores = llm.score(["Question: ... Answer: choice"], logprobs=0)
scores[0]["token_logprobs"]

OpenAI-Compatible Server

Install the optional HTTP stack:

cd /home/lhl/sansho/nano-vllm-sansho
mamba run -n sansho pip install -e '.[server]'

Start a local OpenAI-compatible API server:

CUDA_VISIBLE_DEVICES=2 mamba run -n sansho nanovllm-openai-server \
  --model ../outputs/eval/longrange-hf/swa256_control_500m \
  --served-model-name sansho-swa256-control-500m \
  --host 127.0.0.1 \
  --port 8000 \
  --max-model-len 8192 \
  --max-num-seqs 8

For eval or fixed-shape serving, pre-touch the known Sansho prompt shapes at startup so the first real request does not pay the cold combined-kernel setup cost. The prompt lengths and batch sizes form a cross-product; combinations that exceed max_num_seqs or max_num_batched_tokens are skipped. For graph replay, warm max_tokens - 1 decode steps because the first generated token comes from prefill logits and the graph starts at the first decode step.

CUDA_VISIBLE_DEVICES=0 mamba run -n sansho nanovllm-openai-server \
  --model ../outputs/eval/longrange-hf/swa256_control_500m \
  --max-model-len 8192 \
  --max-num-seqs 8 \
  --max-num-batched-tokens 32768 \
  --sansho-warmup-prompt-lens 512 8064 \
  --sansho-warmup-batch-sizes 1 8 \
  --sansho-warmup-decode-tokens 127 \
  --sansho-decode-graph-replay-max-tokens 128

Supported endpoints:

  • GET /health
  • GET /v1/models
  • POST /v1/completions
  • POST /v1/chat/completions

The server returns OpenAI-compatible logprobs by default because the retained Sansho use case is local evaluation. Disable that default with --disable-default-logprobs, or set "logprobs": false on an individual completion request / "logprobs": false on a chat request. For prompt scoring, use completions with echo=true, max_tokens=0, and logprobs=0 or higher.

The first server cut is non-streaming and serializes generation through one engine lock. Use it for local API smoke tests and lightweight manual testing; use scripts/benchmark_sansho_serving_compare.py from the parent repo for retained performance numbers.

Benchmark

See bench.py for benchmark.

Test Configuration:

  • Hardware: RTX 4070 Laptop (8GB)
  • Model: Qwen3-0.6B
  • Total Requests: 256 sequences
  • Input Length: Randomly sampled between 100–1024 tokens
  • Output Length: Randomly sampled between 100–1024 tokens

Performance Results:

Inference Engine Output Tokens Time (s) Throughput (tokens/s)
vLLM 133,966 98.37 1361.84
Nano-vLLM 133,966 93.41 1434.13

Star History

Star History Chart

About

Nano vLLM

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages