#!/bin/bash
# ─── Qwen SpeedLab: Benchmark Harness ────────────────────
# Comparative benchmarks: Baseline vs MTP vs DFlash vs TurboQuant
# Uses llama-bench for standardized measurements

source "$(dirname "$0")/common.sh"

BENCH_OUT="${LOG_DIR}/benchmark-$(date +%Y%m%d-%H%M%S).txt"

log "═══════════════════════════════════════════════"
log "Qwen SpeedLab Benchmark Suite"
log "Results: ${BENCH_OUT}"
log "═══════════════════════════════════════════════"

{
    echo "=============================================="
    echo "Qwen SpeedLab Benchmarks"
    echo "Date: $(date)"
    echo "GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo 'Unknown')"
    echo "Model: Qwen 3.6 27B Q4_K_M"
    echo "=============================================="
    echo ""

    # ─── 1. Autoregressive Baseline ─────────────────────
    echo "─── Benchmark 1: Autoregressive Baseline ───"
    "${LLAMA_BENCH}" \
        -m "${TARGET_MODEL}" \
        -ngl "${GPU_LAYERS}" \
        -b "${BATCH_SIZE}" \
        -ub "${UBATCH_SIZE}" \
        -p 512 \
        -n 256 \
        -fa 1 \
        --cache-type-k q4_0 \
        --cache-type-v q4_0 \
        2>&1
    echo ""

    # ─── 2. MTP Self-Speculative ────────────────────────
    echo "─── Benchmark 2: MTP Self-Speculative (n_max=3) ───"
    "${LLAMA_BENCH}" \
        -m "${TARGET_MODEL}" \
        -ngl "${GPU_LAYERS}" \
        -b "${BATCH_SIZE}" \
        -ub "${UBATCH_SIZE}" \
        -p 512 \
        -n 256 \
        -fa 1 \
        --cache-type-k q4_0 \
        --cache-type-v q4_0 \
        --spec-type draft-mtp \
        --spec-draft-n-max 3 \
        2>&1
    echo ""

    # ─── 3. MTP with TurboQuant KV ──────────────────────
    echo "─── Benchmark 3: MTP + TurboQuant KV (turbo3) ───"
    "${LLAMA_BENCH}" \
        -m "${TARGET_MODEL}" \
        -ngl "${GPU_LAYERS}" \
        -b "${BATCH_SIZE}" \
        -ub "${UBATCH_SIZE}" \
        -p 512 \
        -n 256 \
        -fa 1 \
        --cache-type-k turbo3 \
        --cache-type-v turbo3 \
        --spec-type draft-mtp \
        --spec-draft-n-max 3 \
        2>&1
    echo ""

    # ─── 4. DFlash Speculative Decoding ─────────────────
    if [ -f "${DFLASH_DRAFT}" ]; then
        echo "─── Benchmark 4: DFlash Speculative Decoding ───"
        "${LLAMA_BENCH}" \
            -m "${TARGET_MODEL}" \
            -md "${DFLASH_DRAFT}" \
            -ngl "${GPU_LAYERS}" \
            -b "${BATCH_SIZE}" \
            -ub "${UBATCH_SIZE}" \
            -p 512 \
            -n 256 \
            -fa 1 \
            --cache-type-k q4_0 \
            --cache-type-v q4_0 \
            --spec-type dflash \
            --spec-dflash-cross-ctx 1024 \
            2>&1
        echo ""
    fi

    # ─── 5. Long context stress test ────────────────────
    echo "─── Benchmark 5: Long Context (32K prompt) ───"
    "${LLAMA_BENCH}" \
        -m "${TARGET_MODEL}" \
        -ngl "${GPU_LAYERS}" \
        -b "${BATCH_SIZE}" \
        -ub "${UBATCH_SIZE}" \
        -p 32768 \
        -n 128 \
        -fa 1 \
        --cache-type-k q4_0 \
        --cache-type-v q4_0 \
        --spec-type draft-mtp \
        --spec-draft-n-max 3 \
        2>&1
    echo ""

} | tee "${BENCH_OUT}"

log "Benchmarks complete! Results saved to ${BENCH_OUT}"

# ─── Summary ─────────────────────────────────────────────
echo ""
echo "═══════════════════════════════════════════════"
echo "Quick Reference: Expected tok/s on RTX 3090"
echo "═══════════════════════════════════════════════"
echo "Baseline (autoregressive):    ~35-45 tok/s"
echo "MTP (self-speculative):       ~55-70 tok/s  (1.5-2x)"
echo "MTP + TurboQuant KV:          ~60-80 tok/s  (1.7-2.2x)"
echo "DFlash (block diffusion):     ~78-130 tok/s (2.2-3.5x)"
echo ""
echo "Note: Speed varies by task type."
echo "Code/structured → highest; prose/creative → lower"