Skip to content

About

Lightweight vLLM-like inference engine written in Rust

Resources

Stars

1 star

Watchers

1 watching

Forks

Latest commit

 

History

114 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

nano-vllm-rs

Overview

nano-vllm-rs is a lightweight vLLM-style inference engine written in Rust with Burn, inspired by GeeeekExplorer/nano-vllm.

The current implementation supports Qwen3-style models from Hugging Face safetensors directories, Qwen3 GGUF files, and Qwen3.6 MoE GGUF models with a CUDA trunk and CPU-resident experts.

Run Inference

Run on CPU:

cargo run --release -- run --device cpu --model ./models/Qwen3-0.6B --prompt "Hello"

Run on CUDA:

cargo run --release --no-default-features --features cuda -- \
  run --device cuda --model ./models/Qwen3-0.6B --prompt "Hello"

Run on ROCm:

cargo run --release --no-default-features --features rocm -- \
  run --device rocm --model ./models/Qwen3-0.6B --prompt "Hello"

Task shortcut:

task run DEVICE=cuda MODEL=./models/Qwen3-0.6B PROMPT="Hello"

Notes:

  • CPU, CUDA, and ROCm use Burn's Dispatch backend.
  • CUDA tasks pin /usr/local/cuda-13.2 by default. Override with CUDA_HOME=....
  • CUDA memory settings can be overridden with NANO_VLLM_MAX_NUM_SEQS, NANO_VLLM_MAX_MODEL_LEN, NANO_VLLM_MAX_BATCHED_TOKENS, NANO_VLLM_GPU_MEMORY_UTILIZATION, and NANO_VLLM_KVCACHE_BLOCK_SIZE.

GGUF Models

GGUF model files can be passed directly as --model.

NANO_VLLM_TOKENIZER_PATH=./models/unsloth-Qwen3-0.6B/tokenizer.json \
cargo run --release --no-default-features --features cuda -- \
  run --device cuda \
  --model ./models/unsloth-Qwen3-0.6B-GGUF/Qwen3-0.6B-Q4_0.gguf \
  --prompt "Hello" \
  --max-tokens 16

Supported GGUF tensor types currently include:

  • F32, F16, BF16
  • Q4_0, Q4_1, Q5_0, Q5_1, Q8_0, Q8_1
  • Q4K, Q5K, Q6K, Q8_0 packed kernels for Qwen3.6 MoE

Qwen3.6 MoE keeps routed expert weights in the GGUF-backed system RAM mapping. Router and shared-expert weights remain on CUDA. Each layer transfers its hidden states and router logits to the host, executes the selected packed experts there, then transfers only the routed hidden output back. Q4_K, Q5_K, and Q6_K experts use Q8_K dynamic activations with AVX2 dispatch on x86_64. Set NANO_VLLM_CPU_THREADS to control the host expert pool; the default is half of the available logical CPUs.

CUDA memory autotuning accounts for resident non-expert weights, paged KV cache, and per-sequence Gated DeltaNet state. Prefix caching is disabled for this hybrid architecture so preempted sequences are recomputed safely.

If the GGUF directory does not contain tokenizer.json, set NANO_VLLM_TOKENIZER_PATH or pass TOKENIZER=... to the Taskfile benchmark task.

Chat REPL

Start an interactive multi-turn chat session with an optional system prompt:

cargo run --release -- repl --device cpu --model ./models/Qwen3-0.6B --system "You are a concise assistant."

Built-in REPL commands:

  • /history: print the current conversation history
  • /clear: clear the conversation history and keep the --system prompt
  • /exit: leave the REPL

The current REPL keeps context by replaying the full transcript each turn. It does not yet persist KV cache across turns.

Benchmark

Use the built-in bench subcommand for Python-compatible benchmark settings:

cargo run --release -- bench \
  --model ./models/Qwen3-0.6B \
  --device cpu \
  --num-seqs 8 \
  --max-input-len 64 \
  --max-output-len 64 \
  --warmup 0 \
  --iters 1 \
  --seed 0 \
  --quiet

Task benchmark entrypoint:

task bench DEVICE=cuda TIER=fast
task bench DEVICE=cuda TIER=normal
task bench DEVICE=cuda TIER=heavy
task bench DEVICE=rocm TIER=fast
task bench DEVICE=cpu TIER=fast

Tiers:

  • fast: NUM_SEQS=8, MAX_MODEL_LEN=128, MAX_INPUT_LEN=64, MAX_OUTPUT_LEN=64, WARMUP=0, ITERS=1
  • normal: NUM_SEQS=16, MAX_MODEL_LEN=512, MAX_INPUT_LEN=256, MAX_OUTPUT_LEN=256, WARMUP=1, ITERS=1
  • heavy: NUM_SEQS=16, MAX_MODEL_LEN=1024, MAX_INPUT_LEN=512, MAX_OUTPUT_LEN=512, WARMUP=1, ITERS=3

Override any tier value on the command line:

task bench DEVICE=cuda TIER=fast NUM_SEQS=1 MAX_MODEL_LEN=32 MAX_INPUT_LEN=8 MAX_OUTPUT_LEN=1

MAX_BATCHED_TOKENS can also be overridden. If omitted, the Taskfile sets it to NUM_SEQS * MAX_INPUT_LEN and clamps it to at least MAX_MODEL_LEN.

Benchmark a GGUF model:

task bench DEVICE=cuda TIER=fast \
  MODEL=./models/unsloth-Qwen3-0.6B-GGUF/Qwen3-0.6B-Q4_0.gguf \
  TOKENIZER=./models/unsloth-Qwen3-0.6B/tokenizer.json

Legacy wrappers are still available:

task bench-cuda-fast
task bench-cuda
task bench-cuda-heavy
task bench-rocm-fast
task bench-rocm
task bench-rocm-heavy

The benchmark prints per-iteration and mean throughput:

  • total generated tokens
  • elapsed time
  • total throughput
  • prefill throughput
  • decode throughput

Recommended benchmark protocol:

  1. Run at least one warmup iteration for stable comparisons, unless using TIER=fast for smoke tests.
  2. Keep MODEL, NUM_SEQS, MAX_INPUT_LEN, and MAX_OUTPUT_LEN fixed when comparing devices or commits.
  3. Compare decode throughput separately from prefill throughput.
  4. GGUF lazy materialization currently avoids full device residency but is much slower than dense safetensors because weights are uploaded as temporary tensors during matmul.

Weight Conversion

Generate a deterministic Qwen3 remap/concatenation plan from Hugging Face safetensors weights:

cargo run --release -- convert-model --model ./models/Qwen3-0.6B --out ./models/Qwen3-0.6B/weights-map.json

About

Lightweight vLLM-like inference engine written in Rust

Resources

Stars

1 star

Watchers

1 watching

Forks

Contributors

Languages