#!/bin/bash
# ─── Qwen SpeedLab: DFlash Mode ──────────────────────────
# Maximum speed via DFlash block diffusion speculative decoding
# Expected: 78-130 tok/s on RTX 3090 (structured output)
# VRAM: ~21-23 GB (target Q4_K_M + draft Q4_K_M + KV cache)

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

log "═══════════════════════════════════════════════"
log "Starting DFlash Mode (Maximum Speed)"
log "═══════════════════════════════════════════════"

# Verify files exist
for f in "${LLAMA_SERVER}" "${TARGET_MODEL}" "${DFLASH_DRAFT}"; do
    if [ ! -f "$f" ]; then
        log "ERROR: Missing file: $f"
        exit 1
    fi
done

# ─── DFlash Configuration ─────────────────────────────────
# Key DFlash params:
#   --spec-type dflash           → Enable DFlash speculative decoding
#   --spec-dflash-cross-ctx N    → Hidden-state context drafter sees (higher = more VRAM)
#   --spec-draft-n-max N         → Max draft tokens (adaptive by default with 'profit' controller)
#   Adaptive draft control       → Automatically adjusts draft depth based on acceptance

exec "${LLAMA_SERVER}" \
    -m "${TARGET_MODEL}" \
    --spec-draft-model "${DFLASH_DRAFT}" \
    --spec-type dflash \
    --spec-dflash-cross-ctx 1024 \
    --port "${PORT}" \
    --host "${HOST}" \
    -np "${N_PARALLEL}" \
    --kv-unified \
    -ngl "${GPU_LAYERS}" \
    --spec-draft-ngl "${GPU_LAYERS}" \
    -b "${BATCH_SIZE}" \
    -ub "${UBATCH_SIZE}" \
    --ctx-size "${CTX_SIZE}" \
    --cache-type-k q4_0 \
    --cache-type-v q4_0 \
    --flash-attn on \
    --jinja \
    --no-mmap \
    --mlock \
    --no-host \
    --reasoning on \
    --chat-template-kwargs '{"preserve_thinking":true}' \
    --temp 0.6 \
    --top-k 20 \
    --top-p 1.0 \
    --min-p 0.0 \
    "$@" \
    2>&1 | tee -a "${LOG_DIR}/dflash-server.log"