開始使用 vLLM TPU

Google Cloud TPU（張量處理單元）可加速機器學習工作負載。vLLM 支援 TPU v6e 和 v5e。有關架構、支援的拓撲結構等資訊，請參閱 TPU 系統架構和特定的 TPU 版本頁面（v5e 和 v6e）。

---

要求

• Google Cloud TPU VM：可以訪問 TPU VM。有關設定說明，請參閱 Cloud TPU 設定指南。
• TPU 版本：v6e, v5e
• Python：3.11 或更新版本（示例中使用 3.12）。

---

安裝

有關使用 pip 安裝 vllm-tpu 或將其作為 Docker 映象執行的詳細步驟，請參閱 安裝指南。

執行 vLLM 伺服器

安裝 vllm-tpu 後，您就可以啟動 API 伺服器了。

1. 登入 Hugging Face：您需要一個 Hugging Face 令牌才能下載模型。

export TOKEN=YOUR_TOKEN
git config --global credential.helper store
huggingface-cli login --token $TOKEN

1. 啟動伺服器：以下命令使用 Llama-3.1-8B 模型啟動伺服器。

vllm serve "meta-llama/Llama-3.1-8B" \
    --download_dir /tmp \
    --disable-log-requests \
    --tensor_parallel_size=1 \
    --max-model-len=2048

1. 傳送請求

伺服器執行後，您可以使用 curl 向其傳送請求。

curl https://:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "meta-llama/Llama-3.1-8B",
        "prompt": "Hello, my name is",
        "max_tokens": 20,
        "temperature": 0.7
    }'

後續步驟：

在 tpu-recipes 倉庫中檢視完整的、端到端的示例配方。

進一步閱讀

• 示例
• 配方
• GKE 使用 vLLM TPU 進行服務