赤兔-Chitu

赤兔 Chitu 是清华大学高性能计算研究所团队与清华系科创企业清程极智联合发布的开源大模型推理引擎框架。“赤兔 Chitu” 引擎的核心亮点在于其能够在非英伟达 Hopper 架构的 GPU 和各类国产芯片上，原生运行 FP8 精度模型。这意味着，无论是依赖于进口芯片还是国产芯片，开发者都能利用这一推理引擎实现高效的 AI 模型推理。这一进展不仅推动了国产 AI 芯片的广泛应用，也为国内的 AI 生态建设提供了强有力的技术支持。

突破与优势

赤兔引擎首次实现了在非英伟达 Hopper 架构 GPU（如旧款 NVIDIA 显卡及国产芯片）上原生运行 FP8 精度模型的突破。传统上，FP8 模型需依赖英伟达最新 H100/H200 显卡，而赤兔通过底层算子的指令级优化（如 GeMM 和 MoE），在不损失模型精度的前提下，将 FP8 数据直接适配到存量硬件上。例如，在 A800 集群的实测中，部署 DeepSeek-671B 满血版推理服务时，赤兔相比主流方案 vLLM 减少 50% 的 GPU 使用量，同时输出速度提升 3.15 倍，显著降低了企业的硬件投入与运营成本。

赤兔的设计注重实际生产需求，支持从单GPU到大规模集群的全场景部署，并针对不同硬件环境提供灵活优化方案。其特点包括：

• 多元算力适配：兼容NVIDIA各代GPU及国产芯片，避免企业对特定硬件的依赖；
• 稳定高效：在低延迟、高吞吐、低显存占用等维度实现平衡，适合长期高并发业务；
• 开源生态共建：通过开源代码吸引开发者与芯片厂商共同优化，加速国产芯片的适配进程。

如何使用

1.从源码安装

git clone --recursive https://github.com/thu-pacman/chitu && cd chitu

pip install -r requirements-build.txt
pip install -U torch --index-url https://download.pytorch.org/whl/cu124  # 根据您的 CUDA 版本调整
TORCH_CUDA_ARCH_LIST=8.6 CHITU_SETUP_JOBS=4 MAX_JOBS=4 pip install --no-build-isolation .

2.单 GPU 推理

torchrun --nproc_per_node 8 test/single_req_test.py \
request.max_new_tokens=64 models=DeepSeek-R1 \
models.ckpt_dir=/data/DeepSeek-R1 infer.pp_size=1 infer.tp_size=8

3.混合并行 (TP+PP)

torchrun --nnodes 2 --nproc_per_node 8 test/single_req_test.py \
request.max_new_tokens=64 infer.pp_size=2 \
infer.tp_size=8 models=DeepSeek-R1 models.ckpt_dir=/data/DeepSeek-R1

4.启动服务

# 在 localhost:21002 启动服务
export WORLD_SIZE=8
torchrun --nnodes 1 \
    --nproc_per_node 8 \
    --master_port=22525 \
    example/serve.py \
    serve.port=21002 \
    infer.stop_with_eos=False \
    infer.cache_type=paged \
    infer.pp_size=1 \
    infer.tp_size=8 \
    models=DeepSeek-R1 \
    models.ckpt_dir=/data/DeepSeek-R1 \
    keep_dtype_in_checkpoint=True \
    infer.mla_absorb=absorb-without-precomp \
    infer.soft_fp8=True \
    infer.do_load=True \
    infer.max_reqs=1 \
    scheduler.prefill_first.num_tasks=100 \
    infer.max_seq_len=4096 \
    request.max_new_tokens=100 \
    infer.use_cuda_graph=True

# 测试服务
curl localhost:21002/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [
      {
        "role": "system",
        "content": "You are a helpful assistant."
      },
      {
        "role": "user",
        "content": "What is machine learning?"
      }
    ]
  }'

5.性能测试

# 使用 benchmark_serving 工具进行全面性能测试
python benchmarks/benchmark_serving.py \
    --model "deepseek-r1" \
    --iterations 10 \
    --seq-len 10 \
    --warmup 3 \
    --base-url http://localhost:21002