目录
Vinci

feat(decoding): GPU sampling and HTTP beam search with radix prefixes (#34)

  • feat(decoding): add GPU top-k/top-p sampling and cached beam search

Keep filtered sampling on CUDA with a startup-allocated CUB workspace and bounded token/score readback. Preserve custom samplers and greedy ABC graphs; use partial selection and in-place weights in the CPU reference sampler.

Add a reusable text BeamSession and standalone worker CLI. Share immutable KV prefixes through reference-counted token slots, fork GDN history through preallocated snapshots, and skip identity copies. Handle EOS, normalized length scores, candidate pruning and session cleanup. HTTP beam scheduling is not part of this entry point. Remove narrow’s temporary shape allocation.

Validate with 179 worker unit/integration tests, 43 core/CPU checks, two CUDA sampler tests, CLI parsing, Qwen3.5-4B beam inference and concurrent HTTP sampling smoke checks. Document usage, supported paths and limitations in docs/DECODING.md.

  • feat(beam): expose HTTP search and share radix prefix ownership

Add non-streaming beam_width and length_penalty to completion and chat requests. Reuse the resident Runtime and KV allocator through incremental, exclusive beam admission, with timeout cancellation and ordinary admission between searches. Reserve candidate and GDN fork workspace at startup.

Move RadixTree into infer-core and fork known chain tips for beam children. Keep shared prefixes pinned, reclaim pruned suffixes, and evict ordinary cached prefixes only when the beam reservation needs the space. Beam trees remain request-local; cross-request GDN prefix caching stays disabled. Bump frontend/worker protocols to reject incompatible deployments.

Validation: CPU/core/protocol/server/scheduler suites and worker hybrid oracle tests pass; CPU Clippy with warnings denied and CUDA release build pass. Qwen3.5-4B on RTX 4070 Ti SUPER verifies both HTTP endpoints, width-one greedy parity, width-four CLI parity, stop/parameter handling, mixed request queuing, chunked prefill, and timeout cancellation followed by healthy reuse.

  • fix(ci): initialize GPU artifact path in runner step
3天前309次提交

RustInfer

A high-performance, architecture-first LLM inference engine written in Rust. OpenAI-compatible API, continuous batching, paged KV cache, and CUDA-graph decode — built on a hexagonal, zero-cost multi-backend core that swaps CUDA for CPU at compile time with no runtime penalty.

License Rust CUDA


Performance

RustInfer outperforms vLLM on an online QPS sweep — Qwen3-4B, NVIDIA H200, max_tokens=512, ignore_eos, matched CUDA-graph decode capture sizes. Across the sweep, RustInfer (red) holds lower TTFT / TPOT / ITL and lower end-to-end latency than vLLM (blue) at equal or higher throughput:

RustInfer vs vLLM — online QPS sweep, H200

RustInfer beats vLLM on the tail, not just the median. Tail inter-token latency (ITL p99) stays below vLLM at every arrival rate — 6.6 → 9.0 ms vs 7.2 → 10.9 ms (qps 1 → 32) — alongside lower ITL / TPOT median and end-to-end latency at matched or higher throughput. Bench harness under bench/.


Design philosophy

RustInfer is organized around a few principles, applied consistently top to bottom.

Hexagonal core (ports & adapters)

infer-core owns nothing but ports — trait definitions for everything the inference path needs from hardware:

infer-core/ports/
  backend.rs      math_ops.rs     fused_ops.rs
  sampler.rs      collective.rs   op_ports.rs

The backends are adapters that implement those ports: infer-backend-cuda (.cu kernels + cuBLASLt + CUTLASS) and infer-backend-cpu (a pure-Rust reference implementation, always linked, used as baseline and for tests). The core has zero knowledge of CUDA; the entire GPU toolchain (nvcc / bindgen / cuDNN / CUTLASS) is confined to the single infer-backend-cuda leaf crate.

Heterogeneous backends at zero cost

The model layer is generic over an LlmBackend trait and monomorphizes at compile time to whichever backend is selected — CUDA or CPU. There is no virtual dispatch on the inference hot path: dispatch cost is paid by the compiler, not per op. The same model code runs on GPU in production and on the CPU reference backend in unit tests, byte-for-byte the same call sites.

High cohesion, low coupling

Eight crates form an acyclic dependency graph with a GPU-free bottom. Each crate has one job; cross-crate contact happens only through infer-protocol (wire types) and infer-core (ports). Swapping a backend, a scheduler policy, or a transport touches exactly one crate.

DDD layering inside the worker

The worker — the most complex crate — is split into Domain / Application / Infrastructure, so pure inference logic never mixes with I/O or orchestration:

infer-worker/src/
  domain/          model.rs, plan.rs, kv, forward_scratch, global_kv_alloc
                   → pure inference logic; no I/O, no transport
  application/     runtime, decode_engine, serve_loop, worker_scheduler,
                   sampler_stack, hosting  → orchestration & lifecycle
  infrastructure/  io, transport           → ZMQ / MsgPack adapters
  components/      attention, ffn, norm, embed, lm_head  → reusable NN blocks
  models/          llama3, qwen3, decoder, loader        → composition

Model variation lives in the data, not in branches

A model’s specialness (quantization, hybrid attention, tied embeddings) is an attribute of the operator/weight it lives on, not a conditional threaded through higher layers. The weight loader is generic and name-driven — it reads exactly the tensor names and shapes the checkpoint declares; only the model module knows how to assemble them.


Architecture

Three cooperating processes share a single TOML config and communicate over ZMQ (IPC) with MessagePack framing:

  infer-server              infer-scheduler                infer-worker
  ┌──────────────┐          ┌──────────────────┐          ┌────────────────────┐
  │ Axum /v1/... │          │ RadixTree prefix   │          │ Runtime<T,D,Model>  │
  │ chat template│  ZMQ     │   cache            │  ZMQ     │  ├ persistent ABC   │
  │ tokenizer    │ ───────► │ continuous batching│ ───────► │  ├ CUDA-graph capture│
  │ SSE stream   │ ◄─────── │ chunked prefill    │ ◄─────── │  └ KV / scratch pool │
  └──────────────┘          └──────────────────┘          │        │            │
         │                          │                       │        ▼            │
         └──────────┬───────────────┘                       │  DecoderModel       │
                    ▼                                        │  (Decoder<T,D>)     │
            infer-protocol                                   │        │            │
  (config / server↔sched / sched↔worker msgs)                │        ▼            │
                                                             │  Components         │
                                                             │ (Attention/FFN/...) │
                                                             └────────┬───────────┘
                                                                      │ calls ops
                                                                      ▼
                                                   infer-core ── LlmBackend (trait port)
                                                                      ▲        ▲
                                                        impl          │        │  impl
                                                  ┌───────────────────┘        └──────────┐
                                          infer-backend-cuda            infer-backend-cpu
                                          (.cu kernels + cuBLASLt)      (reference / tests)

Workspace

Crate Role
infer-core Foundation: dtypes, quant scheme, value types, and the LlmBackend ports. GPU-free — the bottom of the DAG.
infer-protocol Wire types: config parsing + server↔scheduler↔worker messages.
infer-server Axum HTTP front end, OpenAI /v1 API, chat template, SSE streaming.
infer-scheduler RadixTree prefix cache, continuous batching, chunked prefill, batch planning.
infer-worker GPU inference runtime (DDD: domain / application / infrastructure), models, components.
infer-backend-cuda CUDA adapter: .cu kernels + cuBLASLt; statically links the kernel set + CUTLASS.
infer-backend-cpu CPU adapter: pure-Rust reference backend; always linked as baseline / fallback.
infer-frontend Optional front end (outside the core inference path).

Features

  • OpenAI-compatible API/v1/chat/completions and /v1/completions, with SSE streaming, chat templates, and HF tokenizers.
  • Continuous batching with chunked prefill and RadixTree prefix caching.
  • Paged KV cache with profile-driven sizing and KV recycling.
  • CUDA-graph decode — captured graphs over a fixed set of batch sizes, with a persistent ABC buffer that eliminates per-step allocation in the hot loop.
  • Quantization — dense BF16 and AWQ int4 (W4A16) MLP.
  • Models — Llama-3.2, Qwen3, Qwen3 (AWQ), and Qwen3.5-4B text/image inference at TP=1; see Status.

Quick start

Run the prebuilt Docker image (H100 / H200)

The public image contains the compiled CUDA kernels, CUDA 12 runtime, cuBLAS, and cuDNN. The host only needs:

  • an NVIDIA H100 or H200 with a compatible driver;
  • Docker and the NVIDIA Container Toolkit;
  • a local Hugging Face model directory containing config.json, tokenizer.json, and the model weights.

Start the full scheduler + worker + OpenAI-compatible server stack with one command. Replace the host model path and exposed model name as needed:

docker run --rm --gpus all \
  -p 8000:8000 \
  -v /absolute/path/to/Qwen3:/models/model:ro \
  ghcr.io/vinci-hit/rustinfer:1.0.1 \
  serve --model-name Qwen3

Wait for the model to load, then check readiness:

curl --fail http://127.0.0.1:8000/ready

Send a chat completion:

curl http://127.0.0.1:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{"model":"Qwen3",
       "messages":[{"role":"user","content":"What is the capital of France?"}],
       "max_tokens":64}'

No Rust toolchain, CUDA Toolkit, nvcc, libclang, cuDNN headers, or local operator compilation is required. The 1.0.1 image is linux/amd64 and compiled for CUDA architecture sm_90.

Useful container settings include RUSTINFER_MAX_BATCH_TOKENS, RUSTINFER_MAX_BATCH_SEQS, RUSTINFER_MAX_MODEL_LEN, RUSTINFER_CHUNKED_PREFILL_SIZE, and RUST_LOG. A complete custom config can instead be mounted and selected with --config.

Build the Docker image locally

To build from the current source instead of using the published image:

DOCKER_BUILDKIT=1 docker build \
  --build-arg CUDA_ARCH=sm_90 \
  -t rustinfer:local .

Images are architecture-specific. Use CUDA_ARCH=sm_80, sm_86, or sm_89 when building for another supported NVIDIA GPU architecture.

Build from source

Prerequisites

  • Rustup (the repository pins Rust 1.91.1), a CUDA-capable GPU, and the CUDA toolkit.
  • The cuDNN frontend headers on the include path:
export CUDNN_FRONTEND_INCLUDE_DIR=/path/to/site-packages/include

Build

cargo build --release

Run (one-shot e2e smoke test)

Launches scheduler + worker + server for a config, sends one chat completion, prints the reply, and tears everything down:

scripts/e2e_smoke.sh run_qwen3.toml 8100 "Say hello in one short sentence."

The checked-in configs are portable templates. Set their model field to the local Hugging Face model directory before launching; they bind to 127.0.0.1 and use cuda:0 by default.

Run (manual, three processes)

Each binary takes the same --config:

./target/release/rustinfer-scheduler --config run_qwen3.toml &
./target/release/rustinfer-worker    --config run_qwen3.toml &
./target/release/rustinfer-server     --config run_qwen3.toml &

Then hit the OpenAI-compatible endpoint:

curl http://127.0.0.1:8100/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{"model":"Qwen3-4B-Instruct-2507",
       "messages":[{"role":"user","content":"What is the capital of France?"}]}'

Ready-to-use configs: run_qwen3.toml, run_qwen3_awq.toml (AWQ int4), run_llama1b.toml.

Python benchmark tools

The Python project keeps benchmark dependencies optional so building RustInfer does not install another inference engine. Install only the group a script uses:

uv sync --extra bench
uv sync --extra vllm-reference
uv sync --extra sglang-reference

The vLLM and SGLang groups are reference-benchmark environments; they are not runtime dependencies of RustInfer.


Configuration

Config is a single TOML shared by all three processes. Key fields:

Field Meaning
model Path to the HF model directory (config + safetensors + tokenizer).
model_name Name reported by the /v1 API.
device Rank-0 CUDA device, e.g. cuda:0; TP ranks use consecutive devices.
tensor_parallel_size Number of GPUs in the single-process TP group (1 = disabled).
port HTTP port for the server.
tp_operation_timeout_secs Fail-stop deadline for one mirrored TP inference operation.
tp_startup_timeout_secs Longer fail-stop deadline for NCCL and follower startup.
max_batch_tokens Token budget per forward batch.
max_batch_seqs Max concurrent sequences in a batch.
max_model_len Max context length.
chunked_prefill_size Chunked-prefill chunk size (0 = disabled).
enable_prefix_caching Toggle the RadixTree prefix cache.
mem_fraction_static Fraction of GPU memory reserved for weights + static buffers.
num_blocks KV-cache blocks (0 = auto-size from a memory profile).
capture_sizes Batch sizes to capture CUDA graphs for, e.g. [1,2,4,8,16,24,32].
mtp_num_draft_tokens Opt-in Qwen3.5 text MTP; 0 disables it. Requires greedy sampling, TP1, max_batch_seqs=1, and prefix caching disabled. Verification runs eager.
ignore_eos Ignore EOS (useful for fixed-length benchmarking).

Tensor parallelism

Tensor parallelism is disabled by default. To shard a dense BF16 model over two GPUs, set the shared configuration to:

device = "cuda:0"
tensor_parallel_size = 2

One worker process owns the complete TP group. Rank 0 uses device; the other ranks use consecutive CUDA device IDs, so this example uses cuda:0 and cuda:1. The CUDA backend requires NCCL 2.24.3 or newer; the Docker image already includes the matching development and runtime packages.

The implementation supports single-node dense BF16 and block-FP8 Llama/Qwen decoders. Vocabulary size, query/KV head counts, MLP intermediate size, and FP8 weight/scale block boundaries must be evenly divisible by tensor_parallel_size.

CUDA Graph capture is enabled for TP decode and single-sequence prefill. Every rank captures its own device graph and replays the same NCCL collective sequence in lockstep. Mixed prefill+decode batches currently stay eager. TP with AWQ weights, speculative decoding, pipeline parallelism, data parallelism, and expert parallelism are not implemented yet; unsupported combinations fail during startup instead of silently falling back to replicated execution.


Status

RustInfer serves Llama-3.2, Qwen3, and Qwen3-AWQ end to end today.

Qwen3.5-4B runs hybrid Gated DeltaNet/full attention, text and PNG/JPEG data-URL inputs, chunked prefill, and CUDA Graph decode at TP=1. The reproducible service regression covers mixed text/image concurrency and streaming. Hardware/model combinations still require GPU validation; a successful CUDA build alone does not establish inference correctness.

Workers execute a bounded prefill/decode self-check before advertising ready. /health reports HTTP process liveness; /ready requires a loaded Worker group and a fresh scheduler engine heartbeat, and becomes unavailable during failure or shutdown. /metrics exports Prometheus text; /metrics/system retains the browser console’s JSON uptime summary. See metric semantics.

Run the local GPU regression, or provision the automatic master/nightly GPU workflow using the validation guide:

MODEL_PATH=/absolute/path/to/Qwen3.5-4B \
ARTIFACT_DIR=/tmp/rustinfer-gpu-run-001 \
bash scripts/gpu_regression.sh

License

Licensed under the Apache License 2.0. Vendored components retain their own terms; see Third-party notices.

关于

RustInfer是一个用Rust语言实现的高性能大语言模型(LLM)推理引擎,手写CUDA算子,支持BF16与INT4(AWQ)量化推理,单请求decode吞吐量超越vLLM。

44.3 MB
邀请码
    Gitlink(确实开源)
  • 加入我们
  • 官网邮箱:gitlink@ccf.org.cn
  • QQ群
  • QQ群
  • 公众号
  • 公众号

版权所有:中国计算机学会技术支持:开源发展技术委员会
京ICP备13000930号-9 京公网安备 11010802047560号