Auto-tune ctx-size, batch and cache types for low-memory (<=16GB) Spaces; export CACHE_TYPE_{K,V}

start.sh

@@ -13,6 +13,8 @@ export MODEL_PATH="${MODEL_PATH:-/data/models/model.gguf}"
 export CTX_SIZE="${CONTEXT_LENGTH:-${CTX_SIZE:-4096}}"
 export QUANT_PREFERENCE="${QUANT_PREFERENCE:-q4_k_m,q4_k_s,q4_k,q4_0,q4_1}"
 export PERF_PROFILE="${PERF_PROFILE:-balanced}"
+export CACHE_TYPE_K="${CACHE_TYPE_K:-f16}"
+export CACHE_TYPE_V="${CACHE_TYPE_V:-f16}"
 
 # Ensure data directories exist and are writable. Run as root in container startup.
 echo "preparing storage: $HF_HOME and $MODEL_DIR"
@@ -53,6 +55,32 @@ case "$PERF_PROFILE" in
     ;;
 esac
 
+# Adjust defaults for low-memory environments (Spaces free tier: ~16GB)
+if [ -r /proc/meminfo ]; then
+  MEM_KB=$(awk '/MemTotal/ {print $2}' /proc/meminfo || echo 0)
+else
+  MEM_KB=0
+fi
+echo "detected memory (KB): $MEM_KB"
+if [ "$MEM_KB" -gt 0 ] && [ "$MEM_KB" -le 18000000 ]; then
+  echo "low-memory profile detected: adjusting ctx-size and cache types"
+  # reduce default context to save KV cache memory
+  CTX_SIZE="${CTX_SIZE:-2048}"
+  # use compact cache formats
+  CACHE_TYPE_K="${CACHE_TYPE_K:-f16}"
+  CACHE_TYPE_V="${CACHE_TYPE_V:-f16}"
+  # further reduce batch sizes if balanced/throughput
+  if [ "$PERF_PROFILE" = "balanced" ]; then
+    DEFAULT_BATCH=4
+    DEFAULT_UBATCH=4
+  elif [ "$PERF_PROFILE" = "throughput" ]; then
+    DEFAULT_BATCH=64
+    DEFAULT_UBATCH=64
+  fi
+fi
+
+export CACHE_TYPE_K CACHE_TYPE_V CTX_SIZE
+
 server_args=(
   --model "$MODEL_PATH"
   --host 0.0.0.0
@@ -65,8 +93,8 @@ server_args=(
   --batch-size "${BATCH_SIZE:-$DEFAULT_BATCH}"
   --ubatch-size "${UBATCH_SIZE:-$DEFAULT_UBATCH}"
   --ctx-size "${CTX_SIZE:-4096}"
-  --cache-type-k "${CACHE_TYPE_K:-f16}"
-  --cache-type-v "${CACHE_TYPE_V:-f16}"
+  --cache-type-k "${CACHE_TYPE_K}"
+  --cache-type-v "${CACHE_TYPE_V}"
   --metrics
 )