#!/usr/bin/env bash
# k3-test/20_k3_experiments.sh — K3 experiments E0–E8 (adapted to llama.cpp mmap reality).
#
# IMPORTANT FRAMING (see PREP_REPORT.md):
#   UD-Q4_K_XL = 1508.7 GB, box has ~516 GB RAM + 96 GB VRAM. Full residency is
#   IMPOSSIBLE — mmap page-cache streaming IS the deployment mode being tested.
#   "RAM-resident hot set" == warm page cache; "SSD cold tier" == after drop_caches.
#
# Run a subset with: EXP_LIST="E1 E2" ./20_k3_experiments.sh
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/lib/common.sh"

[[ -f $K3_MODEL_PATH ]] || die "K3 not downloaded — run ./01_download_k3.sh (and check config.env K3_SUBDIR)"
CHAT="$K3TEST_ROOT/prompts/chat.txt"; CODING="$K3TEST_ROOT/prompts/coding.txt"; LONG="$K3TEST_ROOT/prompts/longdoc.txt"
EXP_LIST="${EXP_LIST:-E0 E1 E2 E3 E4 E5 E6 E7 E8 E9 E10 E11 E12 E13}"

TSPLIT_HOME="${TSPLIT_HOME:-1,1,1}"   # 3x16GB = 48GB = home 2x3090 VRAM budget
BASE=( -m "$K3_MODEL_PATH" -ngl 999 --tensor-split "$TSPLIT_HOME" --cpu-moe )

have_exp() { [[ " $EXP_LIST " == *" $1 "* ]]; }
save_flagset() { printf '%s\n' "$*" > "$LOG_ROOT/$1.flagset" 2>/dev/null || true; }

# ============================================================================
# E0 — trunk split: 1 vs 2 vs 3 GPUs. UD-Q4_K_XL dense trunk is mostly Q8_0
#      (~55-60 GB) so a single 24 GB card can't hold it; llama.cpp spills the
#      remainder to CPU automatically — that spill cost is exactly what we measure.
# ============================================================================
if have_exp E0; then
  info "=== E0 trunk split sweep ==="
  # HOME CONSTRAINT: 2x3090 = 48 GB VRAM. This box's cards are 16 GB, so:
  #   3 cards (1,1,1)   = 48 GB  == exact home VRAM budget. Trunk ~58 GB Q8_0
  #                        does NOT fit -> ~10 GB trunk spills to RAM (re-read/token).
  #   2 cards (1,1)     = 32 GB  == tighter (e.g. single 3090 + 8 GB elsewhere).
  #   4 cards (1,1,1,1) = 64 GB  == reference "if you added a 3rd 3090-class card".
  drop_caches
  run_exp k3_E0a_home48g  "$CHAT" 128 "${BASE[@]}" --tensor-split 1,1,1     || true
  drop_caches
  run_exp k3_E0b_home32g  "$CHAT" 128 "${BASE[@]}" --tensor-split 1,1       || true
  drop_caches
  run_exp k3_E0c_3card64g "$CHAT" 128 "${BASE[@]}" --tensor-split 1,1,1,1   || true
fi

# ============================================================================
# E1 — decode baseline with hot expert set RAM-resident (warm page cache).
#      Two identical passes: pass A cold (post drop_caches) = streaming,
#      pass B warm = RAM-resident ceiling for this prompt set.
# ============================================================================
if have_exp E1; then
  info "=== E1 cold vs warm decode ==="
  drop_caches
  run_exp k3_E1a_cold "$CODING" 128 "${BASE[@]}"                              || true
  run_exp k3_E1b_warm "$CODING" 128 "${BASE[@]}"                              || true
  run_exp k3_E1c_warm_chat "$CHAT" 128 "${BASE[@]}"                           || true
  # length-graded prompts: disk-read (expert streaming) vs prompt length
  for L in 2048 8192; do
    GP="$K3TEST_ROOT/prompts/gen/coding_${L}tok.txt"
    [[ -f $GP ]] || continue
    run_exp "k3_E1_len${L}" "$GP" 64 "${BASE[@]}"                             || true
  done
fi

# ============================================================================
# E9 — pp scaling (prompt-processing tok/s at 512/2K/8K), warm process,
#      model-load excluded by llama-bench design. Expert-streaming io deltas
#      still visible via the run's cgroup io accounting.
# ============================================================================
if have_exp E9; then
  info "=== E9 pp scaling (llama-bench) ==="
  PP_SIZES="${PP_SIZES:-512 2048 8192}" GEN="${GEN:-32}" \
    "$SCRIPT_DIR/15_pp_scaling.sh" k3_E9 \
    -m "$K3_MODEL_PATH" -ngl 999 --tensor-split "$TSPLIT" --cpu-moe -t "$THREADS" || true
fi

# ============================================================================
# E2/E3 — "prefetch" efficacy, approximated with kernel facilities:
#   E2-on  : mmap + enlarged readahead (layer-sequential access pattern benefits)
#   E3-off : --load-mode dio (direct I/O = every expert read hits NVMe, no cache)
# iostat logs give bytes/token; comparing to E1b isolates SSD-tier cost.
# ============================================================================
if have_exp E2; then
  info "=== E2 mmap + readahead boost ==="
  for f in /sys/block/*/queue/read_ahead_kb; do echo 4096 | $SUDO tee "$f" >/dev/null || echo "CONTAINER_LIMITED: readahead unchanged ($f)" >&2; done
  drop_caches
  run_exp k3_E2_ra4096_cold "$CODING" 128 "${BASE[@]}"                        || true
  for f in /sys/block/*/queue/read_ahead_kb; do echo 128  | $SUDO tee "$f" >/dev/null || true; done
fi

if have_exp E3; then
  info "=== E3 direct-io, no page cache (cold every token) ==="
  drop_caches
  run_exp k3_E3_dio "$CODING" 64 "${BASE[@]}" --load-mode dio                 || true
fi

# ============================================================================
# E4 — cache-size sweep: pin RAM with a /dev/shm filler to shrink page cache,
#      emulating 200/320/450-expert hot sets on the home build's 768 GB.
#      (Rental: ~490 GB usable. Home: ~700 GB usable. Filler emulates smaller.)
#      FILLER_GB_LIST tunes the effective cache: cache ≈ RAM_free_at_start − filler.
# ============================================================================
if have_exp E4; then
  info "=== E4 cache-size sweep (RAM filler) ==="
  for filler in ${FILLER_GB_LIST:-0 128 256}; do
    unfill_ram; [[ $filler -gt 0 ]] && fill_ram "$filler" || true
    drop_caches
    run_exp "k3_E4_filler${filler}g_cold" "$CODING" 96 "${BASE[@]}"           || true
    run_exp "k3_E4_filler${filler}g_warm" "$CODING" 96 "${BASE[@]}"           || true
    unfill_ram
  done
fi

# ============================================================================
# E5 — prompt-lookup (n-gram) speculative decoding. Dynamic cache is written
#      during generation; second pass reads it back for accept-length stats.
# ============================================================================
if have_exp E5; then
  info "=== E5 n-gram speculative — spec-only AND spec+expert-offload (the realistic case) ==="
  # baseline for comparison is E1b (same prompt/config, no spec)
  run_exp k3_E5a_ngram_simple "$CODING" 128 "${BASE[@]}" --spec-type ngram-simple || true
  run_exp k3_E5b_ngram_mapk   "$CODING" 128 "${BASE[@]}" --spec-type ngram-map-k  || true
  run_exp k3_E5c_ngram_cache  "$CODING" 128 "${BASE[@]}" --spec-type ngram-cache  || true
  # NOTE: BASE already includes --cpu-moe (K3 can't fit in VRAM), so every E5 run IS
  # the realistic combined case: experts streaming + spec. E1b is the no-spec baseline
  # with identical offload. (Spec-without-offload is impossible for K3 on this box.)
fi

# ============================================================================
# E10 — K3-trained DSpark drafter (EAGLE-3-class, real MTP descendant) vs n-gram.
#       The actual deployment-grade speculative path. Needs 03_download_draft.sh.
# ============================================================================
if have_exp E10; then
  info "=== E10 DSpark draft-model speculative ==="
  DRAFT="$MODEL_DIR/k3-draft/draft.gguf"
  if [[ -e "$DRAFT" ]]; then
    run_exp k3_E10a_dspark        "$CODING" 128 "${BASE[@]}" --spec-type draft-dspark -md "$DRAFT" || true
    run_exp k3_E10b_dspark_chat   "$CHAT"   128 "${BASE[@]}" --spec-type draft-dspark -md "$DRAFT" || true
    grep -iE "accept|draft" "$LOG_ROOT/k3_E10a_dspark/run.log" | tail -4 || true
  else
    echo "E10 SKIPPED: no draft at $DRAFT (run ./03_download_draft.sh)" | tee "$LOG_ROOT/k3_E10_SKIPPED.txt"
  fi
fi

# ============================================================================
# E6 — GPU hot-expert offload: send experts of the FIRST HOT_LAYERS layers to the
#      GPUs not carrying the trunk (VRAM-permitting). Pattern uses verified tensor
#      names: blk.N.ffn_{up,gate,down}_exps.weight
# ============================================================================
if have_exp E6; then
  info "=== E6 GPU hot-expert offload ==="
  # HOT_EXPERT_RE must match layer indices, e.g. "(0|[1-9]|1[01])" = layers 0–11.
  # Expert tensor names verified in llama.cpp src/llama-arch.cpp:
  #   blk.N.ffn_gate_exps.weight / ffn_up_exps.weight / ffn_down_exps.weight
  HOT_EXPERT_RE="${HOT_EXPERT_RE:-(0|[1-9]|1[01])}"
  # NOTE: with the 48GB home budget the trunk already over-fills VRAM (58>48), so
  # GPU hot-expert offload is NOT available at home unless trunk is requantized to ~Q4.
  # This E6 variant is a REFERENCE (trunk 4 cards + hot experts on CUDA4) to price that.
  # trunk on CUDA0-3 (TSPLIT=1,1,1,1); hot-layer experts to the spare bank CUDA4
  run_exp k3_E6_hotgpu "$CODING" 128 \
      -m "$K3_MODEL_PATH" -ngl 999 --tensor-split 1,1,1,1 \
      -ot "blk.${HOT_EXPERT_RE}.ffn_(up|gate|down)_exps.weight=CUDA4,${EXPERT_CPU_OT}" || true
fi

# ============================================================================
# E7 — long-context probe: 8K prefill then decode; state = 24-layer MLA KV
#      + KDA recurrent state. Reports pp tok/s (prefill cost) and whether decode
#      speed degrades vs E1. Runs at CTX_SIZE from config (default 8192).
# ============================================================================
if have_exp E7; then
  info "=== E7 long-context probe (CTX=$CTX_SIZE) ==="
  drop_caches
  CTX_SIZE="$CTX_SIZE" run_exp k3_E7_longctx "$LONG" 128 "${BASE[@]}"         || true
  grep -E "n_kv|recurrent|state size|llama_kv|memory" "$LOG_ROOT/k3_E7_longctx/run.log" \
      > "$LOG_ROOT/k3_E7_longctx/state_notes.txt" 2>/dev/null || true
fi

# ============================================================================
# E11 — thread-count sweep: SMT (112) vs physical cores (56). Bandwidth-bound MoE
#       GEMM often runs FASTER on physical cores only. Home build = 32C/64T, so the
#       ratio transfers.
# ============================================================================
if have_exp E11; then
  info "=== E11 thread sweep (home-relevant: 3 GPUs, 48GB) ==="
  run_exp k3_E11a_t112 "$CODING" 96 "${BASE[@]}" -t 112 || true
  run_exp k3_E11b_t56  "$CODING" 96 "${BASE[@]}" -t 56  || true
fi

# ============================================================================
# E12 — concurrency: 2 simultaneous decode streams share the page-cached expert
#       reads -> aggregate tok/s is the multi-session answer (plan: 2-4 sessions).
# ============================================================================
if have_exp E12; then
  info "=== E12 two concurrent streams (48GB config) ==="
  start_monitors "$LOG_ROOT/k3_E12_concurrent"
  run_llama k3_E12_s1 "$CODING" 96 "${BASE[@]}" & P1=$!
  run_llama k3_E12_s2 "$CHAT"   96 "${BASE[@]}" & P2=$!
  wait $P1 $P2 || true
  stop_monitors
fi

# ============================================================================
# E13 — DSpark draft tuning: --draft-max (num speculative tokens) sweep.
# ============================================================================
if have_exp E13; then
  info "=== E13 DSpark draft-depth tuning ==="
  DRAFT="$MODEL_DIR/k3-draft/draft.gguf"
  if [[ -e "$DRAFT" ]]; then
    for dm in 3 5 7; do
      run_exp "k3_E13_dmax${dm}" "$CODING" 128 "${BASE[@]}" \
          --spec-type draft-dspark -md "$DRAFT" --draft-max "$dm" || true
    done
  fi
fi

# ============================================================================
# E8 — top-k override 16 -> 12. llama.cpp has no router top-k override flag for
#      kimi-k3 as of LLAMA_REF; probing and logging UNSUPPORTED is the deliverable.
# ============================================================================
if have_exp E8; then
  if "$LLAMA_BIN/llama-cli" --help 2>&1 | grep -qiE "router|top-k.*expert|expert.*top-k"; then
    info "=== E8 top-k override flag exists — see help dump ==="
    "$LLAMA_BIN/llama-cli" --help 2>&1 | grep -iE "router|top-k|expert" > "$LOG_ROOT/k3_E8_flags.txt"
    run_exp k3_E8_topk12 "$CHAT" 128 "${BASE[@]}" --router-top-k 12 || true
  else
    info "=== E8 UNSUPPORTED: no router top-k override in this llama.cpp build — logged ==="
    echo "UNSUPPORTED as of $LLAMA_REF" | tee "$LOG_ROOT/k3_E8_UNSUPPORTED.txt"
  fi
fi

info "=== all requested experiments done. Logs in $LOG_ROOT ==="
ls -dt "$LOG_ROOT"/k3_* 2>/dev/null | head -20