#!/bin/bash
# 部署最终配置：内核版 llama.cpp + MTP 增殖头(精简) + 262K + MTP/ngram 投机解码
set -e
SUDO_PW=yueling408
BINK=/home/zyw/llama.cpp-bonsai-kernel
SRCB=/home/zyw/src/work/llama.cpp-9a9394a/build/bin
MODELD=/home/zyw/models/ternary-bonsai-2-27b

echo "=== [1/5] 二进制归位 ==="
pkill -x llama-server 2>/dev/null || true; sleep 2
mkdir -p "$BINK"
cp -r "$SRCB" "$BINK/"
ls "$BINK/bin/llama-server" && "$BINK/bin/llama-server" --version 2>&1 | head -2

echo "=== [2/5] 模型归位 ==="
mkdir -p "$MODELD"
mv -f /home/zyw/bonsai2-mtp/Ternary-Bonsai-2-27B-PTQ1_0-mtp-lean.gguf "$MODELD/"
ls -la "$MODELD"/*.gguf | awk '{print "  ", $5, $9}'

echo "=== [3/5] 写三个单元 ==="
write_unit() {
  local NAME="$1" DESC="$2" REASON="$3" SAMP="$4"
  cat > /tmp/$NAME.service <<EOF
# /etc/systemd/system/$NAME.service
# 2026-09-22 升级：社区 PTQ1_0 解码内核 + MTP 增殖头（自 Ridge 嫁接）
# 底座 Qwen3.8-27B 三值 PTQ1_0；需自编译的补丁版 fork（见 ~/src/work/llama.cpp-9a9394a）
# $DESC
#
# 相对原版 fork 的实测提升（RTX 5060 Ti 16GB）：
#   llama-bench tg128  45.12 → 56.07 tok/s（+24%，内核）
#   真实任务（内核+MTP+ngram）：抄录 108 / 提取 101 / 创作 63 / 摘要 70 / 数学 71 tok/s
# 262K 满上下文 + MTP 需要"精简版"增殖头（无重复 embedding 表）
[Unit]
Description=$DESC
After=network.target

[Service]
Type=simple
User=zyw
Group=zyw
WorkingDirectory=$MODELD
Environment=LD_LIBRARY_PATH=/usr/local/cuda-13.1/lib64
# 保证小批次数值一致（投机解码无损性的前提）
Environment=GGML_CUDA_BATCH_INVARIANT=1
ExecStart=$BINK/bin/llama-server \\
  -m $MODELD/Ternary-Bonsai-2-27B-PTQ1_0-mtp-lean.gguf \\
  --mmproj $MODELD/Ternary-Bonsai-2-27B-mmproj-Q8_0.gguf \\
  -ngl 99 -fa on -c 262144 -ctk q4_0 -ctv q4_0 -np 1 \\
  --jinja --metrics --warmup --no-context-shift \\
  --spec-type draft-mtp,ngram-map-k4v --spec-draft-n-max 1 \\
  $REASON \\
  $SAMP \\
  --threads 8 --threads-batch 16 \\
  --host 0.0.0.0 --port 8080 --alias bonsai2-27b
Restart=no
TimeoutStartSec=900
TimeoutStopSec=30
StandardOutput=journal
StandardError=journal
SyslogIdentifier=$NAME

[Install]
WantedBy=multi-user.target
EOF
  echo "$SUDO_PW" | sudo -S mv /tmp/$NAME.service /etc/systemd/system/$NAME.service
  echo "$SUDO_PW" | sudo -S chown root:root /etc/systemd/system/$NAME.service
  echo "$SUDO_PW" | sudo -S chmod 644 /etc/systemd/system/$NAME.service
  echo "$SUDO_PW" | sudo -S systemctl disable "$NAME" >/dev/null 2>&1 || true
}

write_unit "llama-bonsai2" \
  "llama.cpp - Bonsai 2 27B [无限制思考 + 内核 + MTP + 262K]" \
  "--reasoning on --reasoning-budget -1" \
  "--temp 1.0 --top-p 0.95 --top-k 20 --repeat-penalty 1.0"

write_unit "llama-bonsai2-think" \
  "llama.cpp - Bonsai 2 27B [思考预算4096 + 内核 + MTP + 262K]" \
  "--reasoning on --reasoning-budget 4096" \
  "--temp 1.0 --top-p 0.95 --top-k 20 --repeat-penalty 1.0"

write_unit "llama-bonsai2-nothink" \
  "llama.cpp - Bonsai 2 27B [关思考+512兜底 + 内核 + MTP + 262K]" \
  "--reasoning off --reasoning-budget 512" \
  "--temp 0.7 --top-p 0.80 --top-k 20 --presence-penalty 1.5 --repeat-penalty 1.0"

echo "$SUDO_PW" | sudo -S systemctl daemon-reload

echo "=== [4/5] 启动并验证 ==="
echo "$SUDO_PW" | sudo -S systemctl reset-failed llama-bonsai2 >/dev/null 2>&1 || true
echo "$SUDO_PW" | sudo -S systemctl start llama-bonsai2
for i in $(seq 1 60); do
  curl -s --noproxy "*" -m 2 http://127.0.0.1:8080/health 2>/dev/null | grep -q ok && { echo "  ✅ 就绪（${i}s）"; break; }
  sleep 1
done
systemctl is-active llama-bonsai2
echo "  显存: $(nvidia-smi --query-gpu=memory.used,memory.total --format=csv,noheader)"
curl -s --noproxy "*" -m 5 http://127.0.0.1:8080/v1/models | python3 -c "
import sys,json; d=json.load(sys.stdin)['data'][0]; m=d.get('meta',{})
print('  模型:',d['id'],'| n_ctx:',m.get('n_ctx'),'| 参数:',m.get('n_params'))
" 2>/dev/null

echo "=== [5/5] 投机解码是否挂上 ==="
echo "$SUDO_PW" | sudo -S journalctl -u llama-bonsai2 -n 60 --no-pager 2>/dev/null | grep -iE "adding speculative|speculative decoding context|n_layer_all" | head -5
echo "=== 完成 ==="
