rust

Rust AI 与大模型集成：Candle、llm.rs 与 Python FFI

Rust AI 开发生态深度解析：HuggingFace Candle 纯 Rust 推理框架、llm.rs 本地大模型部署、ONNX Runtime、Rust 与 Python FFI 互操作、向量数据库 Qdrant 集成，以及 AI 推理性能优化策略。

Leeting Yan 4 分钟阅读 1607 字

目录

  1. 为什么 Rust 适合 AI 推理
  2. Candle：纯 Rust ML 框架
  3. llm.rs 本地大模型推理
  4. ONNX Runtime 集成
  5. Rust ↔ Python FFI 互操作
  6. 向量数据库 Qdrant 集成
  7. RAG 架构实战
  8. 性能优化与部署

1. 为什么 Rust 适合 AI 推理

Python 主导了 AI 训练生态，但 Rust 在 AI 推理和基础设施领域崛起：

维度PythonRust
运行时开销GIL + 解释器开销零成本抽象，直接编译为原生码
内存安全GC + 手动 unsafe（C++ 扩展）编译期保证，无 GC 停顿
部署体积PyTorch ~2GBCandle 模型 + 推理 < 50MB
并发性能GIL 限制真并行，无锁数据结构
跨平台依赖复杂静态链接，单二进制

典型用例：边缘设备推理、API 服务层、RAG 检索、模型量化部署。

2. Candle：纯 Rust ML 框架

HuggingFace Candle 是 HuggingFace 推出的纯 Rust ML 框架，无需 Python 运行时即可运行 Transformer 模型。

2.1 核心设计

use candle_core::{Device, Tensor, DType};

let device = Device::new_cuda(0)?; // 或 Device::Cpu, Device::Metal
let a = Tensor::randn(0f32, 1., (2, 3), &device)?;
let b = Tensor::randn(0f32, 1., (3, 4), &device)?;
let c = a.matmul(&b)?;
println!("{}", c);

2.2 加载预训练模型

use candle_transformers::models::bert::{BertModel, Config};
use tokenizers::Tokenizer;

let device = Device::Cpu;
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_path], DType::F32, &device)? };
let config = Config::v-base();
let model = BertModel::load(vb, &config)?;

let tokenizer = Tokenizer::from_file(tokenizer_path).map_err(E::msg)?;
let tokens = tokenizer.encode("Hello, world!", true).map_err(E::msg)?;
let token_ids = Tensor::new(tokens.get_ids(), &device)?.unsqueeze(0)?;

let embeddings = model.forward(&token_ids)?;
println!("embeddings shape: {:?}", embeddings.shape());

2.3 LLaMA/Mistral 文本生成

use candle_transformers::models::llama::{Llama, Config};

let mut model = Llama::load(vb, &config)?;
let mut logits_processor = LogitsProcessor::new(299792458, None, None);

let mut index_pos = 0;
let mut tokens = prompt_tokens;

for index in 0..sample_len {
    let context_size = if index > 0 { 1 } else { tokens.len() };
    let ctxt = &tokens[tokens.len().saturating_sub(context_size)..];
    let input = Tensor::new(ctxt, &device)?.unsqueeze(0)?;
    let logits = model.forward(&input, index_pos)?;
    let logits = logits.squeeze(0)?.to_dtype(DType::F32)?;
    let next_token = logits_processor.sample(&logits)?;
    tokens.push(next_token);
    index_pos += ctxt.len();
}

2.4 量化模型（GGUF 格式）

use candle_core::quantized::gguf_file;

let mut file = std::fs::File::open("llama-7b-q4_0.gguf")?;
let content = gguf_file::Content::read(&mut file)?;
let model = quantized_llama::Model::from_gguf(content, &mut file, &device)?;
// Q4_0 量化将模型从 13GB 压缩至 ~4GB，推理速度提升 2-3x

3. llm.rs 本地大模型推理

llm.rs（现迁移至 rustformers）支持 GGML/GGUF 格式模型的本地推理。

3.1 基础推理

use llm::models::Llama;
use llm::Model;

let model = llm::load::<Llama>(
    &path,
    llm::TokenizerSource::Embedded,
    Default::default(),
    llm::load_progress_callback_stdout,
)?;

let mut session = model.start_session(Default::default());

session.infer(
    &model,
    &mut rand::thread_rng(),
    &llm::InferenceRequest {
        prompt: "Rust is a programming language".into(),
        ..Default::default()
    },
    &mut Default::default(),
    |token| {
        print!("{}", token);
        Ok(llm::InferenceFeedback::Continue)
    },
)?;

3.2 与 Prompt 模板集成

fn chat_prompt(system: &str, history: &[(String, String)], user: &str) -> String {
    let mut prompt = format!("<|system|>\n{}\n", system);
    for (u, a) in history {
        prompt.push_str(&format!("<|user|>\n{}\n<|assistant|>\n{}\n", u, a));
    }
    prompt.push_str(&format!("<|user|>\n{}\n<|assistant|>\n", user));
    prompt
}

4. ONNX Runtime 集成

ONNX Runtime 支持 Rust 绑定，可运行 PyTorch/TensorFlow 导出为 ONNX 的模型。

4.1 推理示例

use ort::{Environment, SessionBuilder, Value};

let environment = Environment::builder()
    .with_name("my_model")
    .build()?
    .into_arc();

let session = SessionBuilder::new(&environment)?
    .with_model_from_file("model.onnx")?;

// 准备输入
let input = vec![1.0f32; 3 * 224 * 224];
let input_tensor = Value::from_array(([1, 3, 224, 224], input.into_boxed_slice()))?;

// 推理
let outputs = session.run(vec![input_tensor])?;
let predictions: ndarray::ArrayViewD<f32> = outputs[0].try_extract()?;

4.2 与 Candle 对比

特性CandleONNX Runtime
模型来源HuggingFace 生态PyTorch/TensorFlow/sklearn
运行时依赖纯 Rustonnxruntime 系统库
性能良优（Intel/AMD 有加速）
内存占用极低中等
适用场景LLM 推理、边缘部署传统 ML 模型服务

5. Rust ↔ Python FFI 互操作

当必须使用 Python 生态（如 PyTorch 训练结果）时，使用 PyO3 桥接。

5.1 PyO3：Rust 编写 Python 模块

use pyo3::prelude::*;

/// 高性能 Rust 函数暴露给 Python
#[pyfunction]
fn fast_tokenize(text: &str) -> Vec<String> {
    text.split_whitespace()
        .map(|s| s.to_lowercase())
        .collect()
}

#[pymodule]
fn my_rust_tokenizer(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(fast_tokenize, m)?)?;
    Ok(())
}

# Python 调用
import my_rust_tokenizer

tokens = my_rust_tokenizer.fast_tokenize("Hello World FROM Rust")
# ['hello', 'world', 'from', 'rust']

5.2 从 Rust 调用 Python

use pyo3::prelude::*;
use pyo3::types::PyDict;

Python::with_gil(|py| {
    let transformers = py.import("transformers")?;
    let tokenizer = transformers
        .getattr("AutoTokenizer")?
        .call_method("from_pretrained", ("bert-base-chinese",), None)?;
    
    let result = tokenizer.call_method("encode", ("你好世界",), None)?;
    let ids: Vec<u32> = result.extract()?;
    println!("{:?}", ids);
    Ok(())
})

5.3 Maturin：快速构建和发布

# 安装 maturin
pip install maturin

# 开发模式（快速迭代）
maturin develop

# 构建 wheel 分发
maturin build --release

# 发布到 PyPI
maturin publish

6. 向量数据库 Qdrant 集成

6.1 Qdrant Rust 客户端

use qdrant_client::prelude::*;
use qdrant_client::qdrant::{CreateCollection, VectorParams, Distance};

let client = QdrantClient::from_url("http://localhost:6334").build()?;

// 创建集合
client.create_collection(&CreateCollection {
    collection_name: "documents".to_string(),
    vectors_config: Some(VectorParams {
        size: 768,
        distance: Distance::Cosine.into(),
        ..Default::default()
    }.into()),
    ..Default::default()
}).await?;

6.2 嵌入与检索

use qdrant_client::qdrant::{PointStruct, UpsertPointsBuilder};

// 插入向量
let points = vec![PointStruct::new(
    1,                              // point id
    vec![0.05, 0.61, 0.76, /* ... 768 dims */],
    json!({"text": "Rust programming guide"}),
)];

client.upsert_points(UpsertPointsBuilder::new("documents", points)).await?;

// 相似度搜索
let search_result = client.search_points(&SearchPoints {
    collection_name: "documents".to_string(),
    vector: query_embedding,
    limit: 5,
    with_payload: Some(true.into()),
    ..Default::default()
}).await?;

7. RAG 架构实战

检索增强生成（RAG）= 文档分块 → 嵌入 → 向量检索 → 上下文注入生成。

7.1 完整 RAG 流水线

use candle_transformers::models::bert::BertModel;

struct RAGSystem {
    embedder: BertModel,
    qdrant: QdrantClient,
    llm: Llama,
}

impl RAGSystem {
    async fn query(&self, question: &str) -> Result<String> {
        // 1. 将问题转为向量
        let query_embedding = self.embed(question).await?;
        
        // 2. 检索相关文档
        let docs = self.qdrant.search_points(&SearchPoints {
            collection_name: "docs".to_string(),
            vector: query_embedding,
            limit: 3,
            ..Default::default()
        }).await?;
        
        // 3. 构建上下文提示
        let context = docs.result.iter()
            .map(|r| r.payload["text"].as_str().unwrap())
            .collect::<Vec<_>>()
            .join("\n");
        
        let prompt = format!(
            "基于以下上下文回答问题：\n{}\n\n问题：{}\n回答：",
            context, question
        );
        
        // 4. 生成回答
        self.generate(&prompt).await
    }
}

7.2 文档分块策略

fn chunk_text(text: &str, chunk_size: usize, overlap: usize) -> Vec<String> {
    text.chars()
        .collect::<Vec<_>>()
        .windows(chunk_size)
        .step_by(chunk_size - overlap)
        .map(|w| w.iter().collect())
        .collect()
}

8. 性能优化与部署

8.1 模型量化策略

量化级别精度损失体积压缩速度提升
FP32基准1x1x
FP16极小0.5x1.5-2x
INT8小0.25x2-3x
Q4_0 (GGUF)中等0.125x3-5x

8.2 批处理推理

// 批量编码提升 GPU 利用率
let batch_size = 32;
for chunk in inputs.chunks(batch_size) {
    let batch_tensor = Tensor::stack(
        &chunk.iter().map(|t| tokenize(t)).collect::<Result<Vec<_>>>()?,
        0
    )?;
    let outputs = model.forward(&batch_tensor)?;
    // 处理 outputs...
}

8.3 部署方案对比

方案延迟吞吐量复杂度
Candle + axum API低中低
ONNX Runtime + gRPC低高中
vLLM (Python) via FFI极低极高高
边缘设备 (WASM + Candle)-低低

8.4 监控与可观测性

use tracing::{info, instrument};

#[instrument]
async fn inference_handler(request: InferenceRequest) -> Result<InferenceResponse> {
    let start = Instant::now();
    let result = model.generate(&request.prompt).await?;
    info!(
        latency_ms = start.elapsed().as_millis(),
        tokens_generated = result.tokens.len(),
        "inference completed"
    );
    Ok(result)
}

Rust 在 AI 推理领域的生态正在快速成熟：Candle 提供了无需 Python 的纯 Rust 方案，PyO3 桥接现有 Python 资产，ONNX Runtime 兼容传统 ML 模型，Qdrant 向量检索补全了 RAG 架构。对于追求低延迟、小体积、高安全的 AI 基础设施，Rust 是越来越不可替代的选择。

下一篇 → 无

继续阅读

探索更多技术文章

浏览归档，发现更多关于系统设计、工具链和工程实践的内容。

全部文章 返回首页

「rust」更多文章

Rust 错误处理与测试：thiserror、mockall 与属性测试 Rust 错误处理最佳实践：Result/Option 组合子、thiserror 与 anyhow 选型、panic 边界控制、单元/集成/文档测试、mockall 模拟、property-based testing 与代码覆盖率。 Rust 系统编程与性能优化：零拷贝、内存剖析与编译调优 Rust 系统编程深度实践：操作系统原语、零拷贝 I/O、mmap 内存映射、性能剖析（cargo flamegraph）、编译器优化（LTO/PGO）、Benchmark 与内存分析，以及 no_std 嵌入式场景。 Rust 桌面端与 WASM：Tauri、WebAssembly 与嵌入式开发 Rust 跨平台开发全景：Tauri 替代 Electron 的轻量级架构、WASM 编译（wasm-bindgen/wasm-pack）、WASI 运行时、嵌入式 Rust（no_std/embedded-hal），以及跨平台发布策略。