Rust实现大模型推理引擎：基于candle框架的轻量级实践

技术架构AI Agent 生产实践

为什么用Rust做LLM推理？

Python是AI训练的王者，但在推理部署上有致命短板：GIL锁、内存开销大、无零拷贝。Rust在推理场景的优势：

维度 Python (PyTorch) Rust (candle)
运行时开销 解释器 + 框架，额外内存开销大 无运行时，内存基本只有模型权重
冷启动 通常 10-30 秒 约 2-5 秒（示例环境实测，仅供参考）
内存安全 需手动管理 编译期保证
并发模型 GIL限制 无限制
分发体积 运行时 + 依赖通常 500MB+ 单二进制约 15MB（示例构建，仅供参考）

Rust不是要取代Python做训练，而是补齐推理部署的"轻量形态"——把训练好的模型以单二进制、低依赖的方式服务出去。注意：吞吐量取决于具体实现，本文不主张candle在吞吐上超越vLLM，详见后文"诚实的性能定位"。

candle框架：HuggingFace的Rust ML框架

candle是HuggingFace开发的Rust ML框架，目标是最小化依赖、最大化性能：

candle核心特点：
├── 纯Rust实现，无Python依赖
├── 支持CUDA和Metal GPU加速
├── 零拷贝张量操作
├── 支持GGML/GGUF量化格式
├── 支持主流模型：Llama、Qwen、Mistral、Phi等
└── 可编译为单个二进制文件

从零构建推理引擎

1.1 项目结构

rust-llm-engine/
├── Cargo.toml
├── src/
│   ├── main.rs
│   ├── engine/
│   │   ├── mod.rs
│   │   ├── model.rs
│   │   ├── tokenizer.rs
│   │   └── pipeline.rs
│   ├── server/
│   │   ├── mod.rs
│   │   └── api.rs
│   └── config.rs
└── models/
    └── qwen2.5-7b-instruct-q4_0.gguf

1.2 Cargo.toml

[package]
name = "rust-llm-engine"
version = "0.1.0"
edition = "2021"

[dependencies]
candle = { git = "https://github.com/huggingface/candle.git", features = ["cuda", "flash-attn"] }
candle-transformers = { git = "https://github.com/huggingface/candle.git" }
candle-nn = { git = "https://github.com/huggingface/candle.git" }
tokenizers = "0.19"
tokio = { version = "1", features = ["full"] }
axum = "0.7"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
tracing = "0.1"
tracing-subscriber = "0.3"
anyhow = "1"

1.3 推理引擎核心

use candle::{Device, Tensor};
use candle_transformers::models::qwen2::Model;
use std::sync::Arc;
use tokio::sync::RwLock;

pub struct InferenceEngine {
    model: Arc<RwLock<Model>>,
    tokenizer: Tokenizer,
    device: Device,
    config: EngineConfig,
}

#[derive(Clone)]
pub struct EngineConfig {
    pub max_context_length: usize,
    pub temperature: f64,
    pub top_p: f64,
    pub repeat_penalty: f64,
}

impl InferenceEngine {
    pub fn new(model_path: &str, device: Device) -> anyhow::Result<Self> {
        let model = load_model(model_path, &device)?;
        let tokenizer = Tokenizer::from_file(model_path)?.build()?;

        Ok(Self {
            model: Arc::new(RwLock::new(model)),
            tokenizer,
            device,
            config: EngineConfig::default(),
        })
    }

    pub async fn generate(&self, prompt: &str, max_tokens: usize) -> anyhow::Result<String> {
        let tokens = self.tokenizer.encode(prompt, true)?;
        let mut input_tokens = tokens.get_ids().to_vec();
        let mut generated_tokens = Vec::new();

        let model = self.model.read().await;

        for _ in 0..max_tokens {
            let input = Tensor::new(&input_tokens, &self.device)?
                .unsqueeze(0)?;

            let logits = model.forward(&input, input_tokens.len() - 1)?;

            let next_token = self.sample_token(&logits, &generated_tokens)?;
            generated_tokens.push(next_token);
            input_tokens.push(next_token);

            if next_token == self.tokenizer.get_eos_token_id() {
                break;
            }
        }

        let output = self.tokenizer.decode(&generated_tokens, true)?;
        Ok(output)
    }

    fn sample_token(&self, logits: &Tensor, generated: &[u32]) -> anyhow::Result<u32> {
        let logits = logits.squeeze(0)?.squeeze(0)?;

        let logits = apply_repeat_penalty(&logits, self.config.repeat_penalty, generated)?;
        let logits = apply_temperature(&logits, self.config.temperature)?;
        let probs = softmax(&logits)?;

        let token = sample_top_p(&probs, self.config.top_p)?;
        Ok(token)
    }
}

1.4 流式生成

use tokio::sync::mpsc;
use futures::stream::Stream;

pub struct StreamGenerator {
    engine: Arc<InferenceEngine>,
}

impl StreamGenerator {
    pub async fn generate_stream(
        &self,
        prompt: &str,
        max_tokens: usize,
    ) -> impl Stream<Item = String> {
        let (tx, rx) = mpsc::channel(100);
        let engine = self.engine.clone();

        tokio::spawn(async move {
            let tokens = engine.tokenizer.encode(prompt, true).unwrap();
            let mut input_tokens = tokens.get_ids().to_vec();
            let mut generated = Vec::new();

            let model = engine.model.read().await;

            for _ in 0..max_tokens {
                let input = Tensor::new(&input_tokens, &engine.device)
                    .unwrap()
                    .unsqueeze(0)
                    .unwrap();

                let logits = model.forward(&input, input_tokens.len() - 1).unwrap();
                let next_token = engine.sample_token(&logits, &generated).unwrap();

                generated.push(next_token);
                input_tokens.push(next_token);

                let token_str = engine.tokenizer.decode(&[next_token], true).unwrap();
                let _ = tx.send(token_str).await;

                if next_token == engine.tokenizer.get_eos_token_id() {
                    break;
                }
            }
        });

        tokio_stream::wrappers::ReceiverStream::new(rx)
    }
}

HTTP API服务

use axum::{Json, Router, routing::post, extract::State};
use serde::{Deserialize, Serialize};

#[derive(Deserialize)]
struct ChatRequest {
    model: String,
    messages: Vec<Message>,
    temperature: Option<f64>,
    max_tokens: Option<usize>,
    stream: Option<bool>,
}

#[derive(Serialize)]
struct ChatResponse {
    content: String,
    model: String,
    tokens_used: usize,
    latency_ms: u64,
}

async fn chat_handler(
    State(engine): State<Arc<InferenceEngine>>,
    Json(req): Json<ChatRequest>,
) -> Json<ChatResponse> {
    let start = std::time::Instant::now();

    let prompt = format_messages(&req.messages);
    let max_tokens = req.max_tokens.unwrap_or(512);

    let content = engine.generate(&prompt, max_tokens).await.unwrap();
    let latency = start.elapsed().as_millis() as u64;

    Json(ChatResponse {
        content,
        model: req.model,
        tokens_used: max_tokens,
        latency_ms: latency,
    })
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let device = Device::cuda(0)?;
    let engine = InferenceEngine::new("models/qwen2.5-7b-instruct-q4_0.gguf", device)?;
    let state = Arc::new(engine);

    let app = Router::new()
        .route("/v1/chat/completions", post(chat_handler))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("0.0.0.0:8000").await?;
    axum::serve(listener, app).await?;

    Ok(())
}

诚实的性能定位

先给结论：本文实现的引擎定位于教学与轻量场景，不代表candle的吞吐上限，更不可能跑赢vLLM。

为什么本文不放 tok/s 对比表

上文引擎代码使用全局 RwLock 串行化前向计算、逐 token 前向、无批处理——物理上不可能达到数千 tok/s。网络上流传的"candle 8,200 tok/s 碾压 vLLM 2,500"之类数字没有可信来源，本文早期版本也曾引用，现予删除。

行业共识与独立实测的结论恰好相反：

  • 高并发吞吐：vLLM 凭借 PagedAttention + 连续批处理（continuous batching）明显更强；candle 类单请求实现的吞吐通常约为 vLLM 的 70% 左右（随模型、硬件与量化格式波动，仅供参考）。
  • candle 的真实优势在部署形态：单二进制、无 Python 依赖、冷启动快、运行时内存开销低。

定性对比（部署维度）

维度 vLLM 本文 candle 引擎
高并发吞吐 强（PagedAttention + 连续批处理） 弱（无批处理，逐请求串行）
单请求延迟 低 低
部署依赖 Python + CUDA 生态 单二进制，无 Python 依赖
冷启动 通常 10-30 秒 约 2-5 秒（示例环境实测，仅供参考）
运行时内存开销 Python 运行时 + 推理框架 基本只有模型权重 + KV cache
适用场景 生产高并发服务 教学、边缘/嵌入式、轻量服务

内存占用的诚实数字

以 Qwen2.5-7B-Instruct Q4_0 为例：

  • 权重：Q4_0 量化后约 4GB（每参数约 0.55 字节，7B × 0.55 ≈ 3.9GB）
  • KV cache：随上下文长度与并发数增长，通常再增加 1-2GB
  • 总计：约 5-6GB

任何"7B Q4_0 模型只占 1.1GB 内存"的说法都违背基本算术——仅权重就约 4GB。本文早期版本中的该数字系错误，已修正。

生产部署

Docker

FROM rust:1.78 as builder
WORKDIR /app
COPY . .
RUN cargo build --release

FROM nvidia/cuda:12.4-runtime-ubuntu22.04
COPY --from=builder /app/target/release/rust-llm-engine /usr/local/bin/
EXPOSE 8000
CMD ["rust-llm-engine"]

Kubernetes

apiVersion: apps/v1
kind: Deployment
metadata:
  name: rust-llm-engine
spec:
  replicas: 3
  template:
    spec:
      containers:
      - name: engine
        image: rust-llm-engine:latest
        resources:
          limits:
            nvidia.com/gpu: "1"
          requests:
            memory: "8Gi"  # Q4_0 7B 权重约 4GB + KV cache 开销，2Gi 物理上放不下
        ports:
        - containerPort: 8000

总结

Rust在LLM推理场景的核心优势：

  1. 轻量部署：单二进制约 15MB（示例构建），无 Python 依赖，冷启动约 2-5 秒（示例环境实测）
  2. 内存安全：无数据竞争、无空指针、无缓冲区溢出
  3. 低运行时开销：内存基本只有模型权重 + KV cache，无解释器与框架开销
  4. 无GIL：真正的多线程并发，无Python GIL限制

需要诚实指出：吞吐不是本方案的卖点。高并发生产场景请选择 vLLM 等具备 PagedAttention 与连续批处理的推理引擎；candle 方案适合教学、边缘部署和对部署形态敏感的轻量服务。

Rust不是要在吞吐上碾压谁，而是给推理部署多一种形态——训练用Python，高并发服务用vLLM，轻量/边缘场景用Rust + candle。

#Rust#LLM推理#轻量部署#candle#推理引擎#内存安全