海光DCU大模型推理环境全流程脚本集｜纯代码可直接复用

整理一套从环境校验、依赖安装、服务部署到压测验证的全链路可运行脚本，全程无额外依赖，复制到服务器按顺序执行即可完成整套推理环境搭建，适合快速部署生产级推理服务。

1. DTK环境一键校验脚本

执行前先确认DTK环境是否正常、硬件是否可识别，避免后续安装全做无用功。
保存为check_dtk_env.sh：

#!/bin/bash
echo "=== DCU 环境基线校验 ==="

# 1. 检查rocm-smi工具
echo -n "1. rocm-smi 工具: "
if command -v rocm-smi &>/dev/null; then
    echo "正常"
else
    echo "未找到，请先加载DTK环境变量"
    exit 1
fi

# 2. 检查HIP编译器
echo -n "2. hipcc 编译器: "
if command -v hipcc &>/dev/null; then
    hipcc --version 2>/dev/null | grep "HIP version"
else
    echo "未找到"
    exit 1
fi

# 3. 设备识别检查
echo -n "3. DCU 设备识别: "
card_count=$(rocm-smi --showhw 2>/dev/null | grep -c "Device ID")
if [ "$card_count" -gt 0 ]; then
    echo "检测到 $card_count 张卡"
else
    echo "未识别到设备，请检查驱动与权限"
    exit 1
fi

# 4. 输出核心硬件信息
echo -e "\n=== 硬件明细 ==="
rocm-smi --showhw 2>/dev/null | grep -E "Device ID|VRAM Total|Compute Unit"

echo -e "\n✅ 环境校验通过"

使用方式：bash check_dtk_env.sh，任意一项失败都会直接提示对应问题。

2. PyTorch与推理依赖一键安装

自动匹配DTK版本安装对应ROCm版PyTorch，同时安装vLLM、transformers等推理必备依赖。
保存为install_deps.sh：

#!/bin/bash
echo "=== 安装 DCU 推理依赖 ==="

# 自动识别DTK对应ROCm版本
DTK_VER=$(hipcc --version 2>/dev/null | grep -oP "HIP version: \K[0-9]+" | head -1)
if [ "$DTK_VER" -ge 26 ]; then
    ROCM_VER="6.0"
    TORCH_VER="2.4.0"
elif [ "$DTK_VER" -ge 25 ]; then
    ROCM_VER="5.7"
    TORCH_VER="2.3.1"
else
    echo "DTK版本过低，建议升级至25.04以上"
    exit 1
fi
echo "匹配 ROCm 版本: $ROCM_VER，PyTorch 版本: $TORCH_VER"

# 安装PyTorch
pip install torch==$TORCH_VER torchvision==0.19.0 \
  --index-url https://download.pytorch.org/whl/rocm$ROCM_VER

# 安装推理常用依赖
pip install transformers accelerate sentencepiece \
  modelscope vllm==0.18.1 requests

# 验证安装结果
echo -e "\n=== 安装结果验证 ==="
python3 << 'EOF'
import torch
import transformers
import vllm
print(f"PyTorch: {torch.__version__}")
print(f"Transformers: {transformers.__version__}")
print(f"vLLM: {vllm.__version__}")
print(f"DCU可用: {torch.cuda.is_available()}")
print(f"DCU数量: {torch.cuda.device_count()}")
EOF

echo "✅ 依赖安装完成"

使用方式：bash install_deps.sh，自动适配DTK版本，无需手动对应。

3. vLLM推理服务一键启停脚本

带进程守护、日志管理的服务化脚本，支持启动、停止、状态查询，无需手动管理进程。
保存为vllm_service.sh：

#!/bin/bash
# 配置项 - 根据实际环境修改
MODEL_PATH="./models/Qwen2-7B-Instruct"
MODEL_NAME="Qwen2-7B-Instruct"
PORT=8000
TP_SIZE=1
DTK_ENV="/opt/dtk-26.04/env.sh"
LOG_FILE="./vllm_service.log"
PID_FILE="./vllm.pid"

# 加载DTK环境
source $DTK_ENV
export TRITON_HIP_LLD_PATH=/opt/dtk-26.04/llvm/bin/ld.lld

start_service() {
    if [ -f "$PID_FILE" ] && kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then
        echo "服务已在运行，PID: $(cat "$PID_FILE")"
        exit 1
    fi
    
    echo "启动 vLLM 推理服务..."
    nohup numactl --cpunodebind=0 --membind=0 \
        vllm serve "$MODEL_PATH" \
        --served-model-name "$MODEL_NAME" \
        --tensor-parallel-size "$TP_SIZE" \
        --dtype bfloat16 \
        --max-model-len 8192 \
        --gpu-memory-utilization 0.9 \
        --enforce-eager \
        --max-num-seqs 128 \
        --host 0.0.0.0 \
        --port "$PORT" \
        >> "$LOG_FILE" 2>&1 &
    
    echo $! > "$PID_FILE"
    echo "服务已启动，PID: $!"
    echo "日志文件: $LOG_FILE"
    echo "端口: $PORT"
}

stop_service() {
    if [ ! -f "$PID_FILE" ]; then
        echo "服务未运行"
        exit 1
    fi
    pid=$(cat "$PID_FILE")
    if kill -0 "$pid" 2>/dev/null; then
        echo "停止服务，PID: $pid"
        kill "$pid"
        sleep 3
        kill -9 "$pid" 2>/dev/null
    fi
    rm -f "$PID_FILE"
    echo "服务已停止"
}

status_service() {
    if [ -f "$PID_FILE" ] && kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then
        echo "✅ 服务运行中，PID: $(cat "$PID_FILE")"
        echo "端口: $PORT"
    else
        echo "❌ 服务未运行"
    fi
}

case "$1" in
    start) start_service ;;
    stop) stop_service ;;
    restart) stop_service; sleep 2; start_service ;;
    status) status_service ;;
    *) echo "用法: $0 start|stop|restart|status"; exit 1 ;;
esac

使用方式：

bash vllm_service.sh start   # 启动服务
bash vllm_service.sh status  # 查看状态
bash vllm_service.sh stop   # 停止服务

4. 推理接口压测脚本

服务启动后，批量并发压测，输出吞吐、平均延迟等核心指标，量化验证服务性能。
保存为benchmark.py：

import time
import requests
import threading
from statistics import mean

# 配置
API_URL = "http://localhost:8000/v1/chat/completions"
MODEL_NAME = "Qwen2-7B-Instruct"
CONCURRENT = 16  # 并发数
TOTAL_REQUESTS = 64
MAX_TOKENS = 256
PROMPT = "请详细介绍海光DCU的主要应用场景和技术特点"

# 全局统计
success = 0
total_tokens = 0
latencies = []
lock = threading.Lock()

def worker():
    global success, total_tokens
    payload = {
        "model": MODEL_NAME,
        "messages": [{"role": "user", "content": PROMPT}],
        "max_tokens": MAX_TOKENS,
        "temperature": 0.1,
        "stream": False
    }
    start = time.perf_counter()
    try:
        resp = requests.post(API_URL, json=payload, timeout=60)
        cost = time.perf_counter() - start
        if resp.status_code == 200:
            data = resp.json()
            tokens = data["usage"]["completion_tokens"]
            with lock:
                success += 1
                total_tokens += tokens
                latencies.append(cost)
    except Exception as e:
        print(f"请求失败: {e}")

if __name__ == "__main__":
    print(f"开始压测: 并发={CONCURRENT}, 总请求={TOTAL_REQUESTS}")
    start_time = time.perf_counter()
    
    threads = []
    for _ in range(TOTAL_REQUESTS):
        t = threading.Thread(target=worker)
        threads.append(t)
        t.start()
        # 控制并发数
        while threading.active_count() > CONCURRENT + 1:
            time.sleep(0.01)
    
    for t in threads:
        t.join()
    
    total_time = time.perf_counter() - start_time
    print("\n=== 压测结果 ===")
    print(f"成功请求: {success}/{TOTAL_REQUESTS}")
    print(f"总生成Token: {total_tokens}")
    print(f"总耗时: {total_time:.2f}s")
    print(f"平均吞吐: {total_tokens / total_time:.2f} token/s")
    print(f"单请求平均延迟: {mean(latencies):.2f}s")
    print(f"QPS: {success / total_time:.2f} req/s")

使用方式：python3 benchmark.py，自动输出完整性能指标。

5. 显存与进程一键清理脚本

服务异常退出、显存残留时执行，快速释放所有DCU资源，无需重启节点。
保存为clean_dcu.sh：

#!/bin/bash
echo "=== DCU 资源清理工具 ==="

# 列出所有占用DCU的进程
echo "当前占用DCU设备的进程:"
echo "PID    用户    进程名"
echo "-------------------------"

pids=""
for dev in /dev/dri/renderD*; do
    dev_pids=$(fuser "$dev" 2>/dev/null)
    for pid in $dev_pids; do
        if ! echo "$pids" | grep -q "$pid"; then
            pids="$pids $pid"
            user=$(ps -p "$pid" -o user= 2>/dev/null)
            comm=$(ps -p "$pid" -o comm= 2>/dev/null)
            printf "%-6s %-8s %s\n" "$pid" "$user" "$comm"
        fi
    done
done

if [ -z "$pids" ]; then
    echo "无占用进程，显存空闲"
    exit 0
fi

read -p "是否终止以上所有进程并释放显存? (y/n): " confirm
if [ "$confirm" = "y" ]; then
    for pid in $pids; do
        kill -9 "$pid" 2>/dev/null && echo "已终止进程 $pid"
    done
    # 清理PyTorch缓存
    python3 -c "import torch; torch.cuda.empty_cache()" 2>/dev/null
    echo "✅ 清理完成"
else
    echo "已取消操作"
fi

使用方式：bash clean_dcu.sh，先列进程再确认，避免误杀正常任务。

以上5个脚本覆盖了DCU大模型推理环境从部署到运维的全流程高频场景，全部基于原生命令与官方工具实现，无额外依赖，存到服务器固定目录即可长期复用。