Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

History

656 Commits

Repository files navigation

PegaInfer logo

PegaInfer

Pure Rust + CUDA LLM inference engine. No PyTorch. No model framework runtime.

Docs & Blog at open-infer.orgJoin the PegaInfer Slack

Quickstart · Models · API · Performance · Architecture · Blog


PegaInfer is an LLM inference engine built entirely in Rust and CUDA — no PyTorch, no ONNX, no framework runtime, every kernel and scheduler hand-written.

It serves frontier-scale models, from Qwen3 to the trillion-parameter Kimi-K2, and already holds its own against the best open-source inference frameworks.

Docs, guides, and engineering deep-dives live at open-infer.org — start with PegaInfer 0.1.0: Writing a Production-Grade Inference Engine in Rust and Co-locating Prefill and Decode on One GPU.

Quickstart

Prerequisites

  • Rust (2024 edition), CUDA Toolkit (nvcc, cuBLAS), CUDA-capable GPU
  • NVIDIA driver R545 (CUDA 12.3) or newer; cuFuncGetName sets this floor, while per-symbol lazy loading keeps the cuda-12090 cudarc binding from requiring a CUDA 12.9 driver
  • The default build (Qwen3-4B / 8B) is pure Rust + CUDA — no Python at all
  • Python 3 + Triton for qwen35 feature builds (build-time only — no Python at runtime)
  • The kimi-k2 EP path additionally needs NCCL ≥ 2.27 at runtime (ncclAlltoAll)

Build & Run

# Download a model
huggingface-cli download Qwen/Qwen3-4B --local-dir models/Qwen3-4B
# Build & start server on port 8000 — no Python needed for the default Qwen3 buildexport CUDA_HOME=/usr/local/cuda
cargo run --release

Note: The server CLI is in pegainfer-server. Model crates such as pegainfer-qwen3, pegainfer-qwen35, and pegainfer-kimi-k2 contain model logic and diagnostics but are not server entrypoints. Use cargo run --release from the workspace root, or cargo run --release -p pegainfer-server -- --model-path <path>.

# Try it
curl -s http://localhost:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{"prompt": "The capital of France is", "max_tokens": 32}'# Streaming
curl -N http://localhost:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{"prompt": "Write a haiku about Rust:", "max_tokens": 64, "stream": true}'

Always use --release. Debug builds are extremely slow for GPU/CUDA code.

More options
# Qwen3.5 requires the feature-gated Triton AOT kernels (Python + Triton at build time)
uv venv && uv pip install triton
export PEGAINFER_TRITON_PYTHON=.venv/bin/python
cargo run --release --features qwen35 -- --model-path models/Qwen3.5-4B
# Disable CUDA Graph (useful for debugging)
cargo run --release -- --cuda-graph=false

Environment variables:

VariableDescription
CUDA_HOMECUDA Toolkit path (default: /usr/local/cuda)
PEGAINFER_TRITON_PYTHONPython with Triton for qwen35 build-time AOT compilation
PEGAINFER_TILELANG_PYTHONPython with TileLang for the glm52 sparse-MLA build-time kernel generation (sm_90a)
PEGAINFER_CUDA_SMGPU SM target override when nvidia-smi unavailable (e.g. 120)
Windows
$env:CUDA_PATH="C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.x"# Default Qwen3 build needs no Python
cargo build --release
cargo run --release -p pegainfer-server ----model-path models/Qwen3-4B
# Qwen3.5 additionally needs Triton for the feature-gated AOT kernels
uv venv .venv --python 3.12
uv pip install "triton-windows<3.7"$env:PEGAINFER_TRITON_PYTHON=".venv\Scripts\python.exe"
cargo run --release --features qwen35 ----model-path models/Qwen3.5-4B

Supported Models

ModelArchitectureParamsStatus
Qwen3-4BFull attention (GQA)4BGreedy + sampling, default feature, pure Rust + CUDA build
Qwen3-8BFull attention (GQA)8BGreedy + sampling, default feature, pure Rust + CUDA build
Qwen3.5-4B / 9B / 27BHybrid Gated DeltaNet + full attention4B / 9B / 27BText-only BF16, greedy + sampling, feature-gated, --features qwen35 (build-time Triton)
DeepSeek-V2-LiteMoE + EP15.7B total / 2.4B activeFeature-gated, --features deepseek-v2-lite, 2-GPU EP2 correctness path
Kimi-K2-InstructMLA + MoE + Marlin INT41T total / 32B activeFeature-gated, --features kimi-k2, 8-GPU EP path

Model type is auto-detected from config.json — just point --model-path at any supported model directory. Every model line is controlled by a cargo feature; only qwen3 is on by default, so the stock build serves Qwen3 with zero Python. Other lines require rebuilding pegainfer-server with the matching --features ... flag before launch.

DeepSeek support is intentionally narrower than the Qwen paths:

  • DeepSeek-V2-Lite requires --features deepseek-v2-lite and the 2-GPU EP2 path. Correctness, direct decode diagnostics, and retained HTTP SLO reports use separate entry points and claim boundaries; see benchmarking.md, status.md, and hf-accuracy-gate.md.

API

OpenAI-compatible /v1/completions endpoint.

FieldTypeDefaultDescription
promptstring(required)Input text
max_tokensint128Maximum tokens to generate
temperaturefloat0.0Sampling temperature (0 = greedy)
top_kint50Top-k sampling
top_pfloat1.0Nucleus sampling threshold
streamboolfalseEnable SSE streaming

Sampling and logprob support is model-dependent; Qwen models support the sampling controls above.

Performance

Single RTX 5090 (32 GB), Qwen3-4B, BF16, TP1 — PegaInfer @ 0b42ed3, vLLM 0.22.1, both driven by the same vllm bench serve client (prefix cache on, seed 42, 1k-in / 128-out). Full tables and method are in the benchmark report; the story behind these numbers is in the 0.1.0 release blog.

Footprint

No framework runtime means a small process that starts fast — one process, no torch.compile:

MetricPegaInfervLLM 0.22.1
Resident memory (idle, loaded)771 MB3814 MB
Startup → HTTP-ready (cold)3.0 s70.0 s
Startup (warm compile cache)~3.0 s32.7 s

~5× smaller resident footprint, and a 3 s cold start against vLLM's 70 s — still 11× even versus vLLM's warm torch.compile cache. PegaInfer is a single process; vLLM's RSS is summed across its process tree.

Qwen3-4B on one RTX 5090: output throughput vs request rate, and warm-cache TTFT vs input length — PegaInfer vs vLLM 0.22.1

Under serving load

Poisson arrivals, 1k-token prompts, 128-token outputs. Throughput tracks vLLM step-for-step through the knee and edges ahead at saturation (1794 vs 1692 tok/s, ~14.0 vs 13.2 req/s at QPS 16). vLLM keeps a per-token decode (TPOT) edge at mid load (QPS 8–12); both knee around QPS 10–12, past which the queue dominates. The saturated-throughput cap from the earlier run is gone — batched lm_head + sampling (#362) lifted it.

Warm-cache latency — the chat / agent hot path

On the multi-turn chat and agent hot path, most of the prompt lands as a warm prefix-cache hit. PegaInfer's first token stays flat as context grows — ~9 ms at 1k tokens, ~26 ms at 16k against vLLM's ~96 ms (3.6×) — with p99 within ~1 ms of p50 at every length. Cold (uncached) prefill is at parity (~1.1 s at 16k).

KV offload — host-tier restore (pegaflow)

With --kv-offload, prefixes evicted from HBM are restored from host DRAM instead of recomputed. At 16k that turns a 1.14 s cold prefill into a 126 ms host-tier restore (9.1×; 2.6× at 256 tokens). The tiering ladder at 16k: HBM hit ~26 ms < host-tier restore ~126 ms ≪ cold prefill ~1.14 s.

Qwen3.5-4B vs current vLLM

Single RTX 5090 (32 GB), Qwen3.5-4B, BF16, TP1 — PegaInfer with the Qwen3.5 decode-tuning change, vLLM 0.23.0, both driven by vllm bench serve 0.23.0. Fixed random prompts, 64 measured requests, 2 warmups, text-only serving with prefix cache off on both engines. Full flags and caveats are in the Qwen3.5 benchmark report.

WorkloadMetricPegaInfervLLM 0.23.0
1 input / 256 outputTPOT mean6.282 ms6.214 ms
1 input / 512 outputTPOT mean6.381 ms6.221 ms
1024 input / 256 outputreported input tokens63,459 (992/request)65,536 (1,024/request)
1024 input / 256 outputTTFT mean (client-contract)55.3 ms66.3 ms
1024 input / 256 outputTPOT mean7.110 ms6.346 ms
1024 input / 256 outputoutput tok/s137.0151.9
2048 input / 1 outputreported input tokens126,957 (1,984/request)131,072 (2,048/request)
2048 input / 1 outputTTFT mean (client-contract)97.4 ms101.9 ms

The decode-tuning change improves PegaInfer's own direct Qwen3.5 decode TPOT by about 2-3%. Against vLLM, prompt-len-1 decode is close, but vLLM still leads the 1024/256 decode and high-concurrency HTTP rows. TTFT rows are fixed-client timings because reported prompt-token totals differ on the longer prompts.

Architecture

flowchart TB
api["HTTP / OpenAI-compatible /v1/completions"]
frontend["pegainfer-server<br/>pegainfer-vllm-frontend"]
runtime["EngineHandle / GenerateRequest / TokenEvent<br/>pegainfer-engine contract · pegainfer-core runtime"]
api --> frontend
frontend --> runtime
subgraph engines["Per-model engine crates"]
direction LR
qwen3["pegainfer-qwen3<br/>full attention"]
qwen35["pegainfer-qwen35<br/>24 linear + 8 full attention"]
dsv2["pegainfer-deepseek-v2-lite<br/>MoE + EP"]
kimi["pegainfer-kimi-k2<br/>MLA + MoE + Marlin INT4"]
end
runtime --> qwen3
runtime --> qwen35
runtime --> dsv2
runtime --> kimi
subgraph shared["Shared kernels and KV management"]
direction LR
kernels["pegainfer-kernels"]
kvcache["pegainfer-kv-cache<br/>pegainfer-kv-offload"]
kvbm["kvbm-logical<br/>ported from NVIDIA Dynamo"]
end
qwen3 --> kernels
qwen35 --> kernels
dsv2 --> kernels
kimi --> kernels
qwen3 --> kvcache
kimi --> kvcache
kvcache --> kvbm
subgraph backends["Backend libraries and communication"]
direction LR
cuda["CUDA"]
cublas["cuBLAS"]
triton["Triton AOT"]
tilelang["TileLang"]
flashinfer["FlashInfer"]
nccl["NCCL"]
deepep["DeepEP shim<br/>NCCL"]
end
kernels --> cuda
kernels --> cublas
kernels --> triton
kernels --> tilelang
kernels --> flashinfer
dsv2 --> nccl
kimi --> deepep
deepep --> nccl
Loading

Key design decisions:

  • GPU-first runtime — model execution stays in native Rust/CUDA paths
  • Custom GPU kernels — CUDA for decode-critical paths, Triton AOT for Qwen3.5 compatibility kernels, FlashInfer for paged attention/sampling, NCCL for multi-GPU reductions, and cuBLAS for matrix multiplication
  • CUDA Graph on Qwen decode paths — eliminates kernel launch overhead where enabled
  • Per-model crate boundary — Qwen3-4B owns its config, weights, scheduler/executor, tests, benches, and kernel plan in pegainfer-qwen3

Model details:

  • Qwen3: 32 Q heads, 8 KV heads (GQA 4:1), head_dim=128
  • Qwen3.5: hybrid — 24 linear attention layers (Gated Delta Rule) + 8 full attention layers, head_dim=256
  • DeepSeek V2-Lite: feature-gated 2-GPU EP2 correctness/attribution path for the HF/host-staged/NCCL narrow greedy gate

What's not (yet) implemented

  • General-purpose quantization for the Qwen lines — INT4 and FP8 today are model-specific (Kimi-K2 Marlin INT4, GLM5.2 FP8), not yet available for the BF16 Qwen models

Development

Fresh-box dev setup

scripts/setup_dev.sh bootstraps a build environment on any fresh NVIDIA Ubuntu host: apt build deps + protobuf-compiler, uv, the rustup nightly pinned by rust-toolchain.toml, the vendored flashinfer/3rdparty/cccl submodule, then cargo build --release. CUDA is a prerequisite — it detects nvcc and fails loudly rather than installing a toolkit, so boot a CUDA image.

bash scripts/setup_dev.sh
# on a GPU whose arch the kernels don't target (e.g. V100 sm_70), compile for another:
PEGAINFER_CUDA_SM=90 bash scripts/setup_dev.sh

To get the box itself, scripts/prime_devbox.sh provisions the cheapest match on Prime Intellect, has the box git-clone this repo over HTTPS, and runs setup_dev.sh — see the script header for one-time setup.

Tests

# Unit tests
cargo test --release --workspace --lib
# Accuracy and integration tests (need GPU + model weights)
PEGAINFER_TEST_MODEL_PATH=models/Qwen3-4B cargo test --release -p pegainfer-qwen3 --test hf_golden_gate
PEGAINFER_TEST_MODEL_PATH=models/Qwen3.5-4B cargo test --release -p pegainfer-qwen35 --features qwen35 --test hf_golden_gate
PEGAINFER_TEST_MODEL_PATH=models/Qwen3.5-4B cargo test --release -p pegainfer-qwen35 --features qwen35 --test e2e_scheduler
PEGAINFER_TEST_MODEL_PATH=models/DeepSeek-V2-Lite cargo test --release -p pegainfer-deepseek-v2-lite --features deepseek-v2-lite --test e2e_ep2 -- --nocapture

The DeepSeek-V2-Lite E2E is a correctness/integration gate. Direct diagnostics and HTTP SLO report commands live in benchmarking.md.

License

Apache-2.0 — see LICENSE and NOTICE. Components ported from NVIDIA Dynamo (the kvbm/kvbm-logical crate) retain their original Apache-2.0 headers; see NOTICE_DYNAMO.

Star History

Star History Chart

Releases

Packages

Contributors

Languages