Add QUANT_PREFERENCE and PERF_PROFILE (low_latency/balanced/throughput) with server defaults

download_model.py

@@ -16,33 +16,31 @@ def require_env(name: str) -> str:
 
 def quant_rank(file_name: str) -> tuple[int, str]:
     name = file_name.lower()
-
-    q8_tokens = ("q8_k_l", "q8_k_m", "q8_k_s", "q8_k", "q8_0", "q8_1")
-    if any(token in name for token in q8_tokens):
-        return (0, file_name)
-
-    if "f16" in name or "bf16" in name:
-        return (1, file_name)
-
-    if any(token in name for token in ("q6_k", "q6_0", "q6_1", "q6")):
-        return (2, file_name)
-
-    if any(token in name for token in ("q5_k_m", "q5_k_s", "q5_k", "q5_0", "q5_1", "q5")):
-        return (3, file_name)
-
-    if any(token in name for token in ("q4_k_m", "q4_k_s", "q4_k", "q4_0", "q4_1", "q4")):
-        return (4, file_name)
-
-    if any(token in name for token in ("q3_k_m", "q3_k_s", "q3_k", "q3_0", "q3_1", "q3")):
-        return (5, file_name)
-
-    if any(token in name for token in ("q2_k", "q2_0", "q2_1", "q2")):
-        return (6, file_name)
-
-    if "f32" in name:
-        return (7, file_name)
-
-    return (8, file_name)
+    # Allow user-specified preference list via QUANT_PREFERENCE env (comma-separated)
+    pref = os.environ.get("QUANT_PREFERENCE", "q4_k_m").split(",")
+    pref = [p.strip().lower() for p in pref if p.strip()]
+    for idx, token in enumerate(pref):
+        if token in name:
+            return (idx, file_name)
+
+    # fallback groups (lower priority than explicit preference)
+    groups = [
+        ("q8_k_l", "q8_k_m", "q8_k_s", "q8_k", "q8_0", "q8_1"),
+        ("f16", "bf16"),
+        ("q6_k", "q6_0", "q6_1", "q6"),
+        ("q5_k_m", "q5_k_s", "q5_k", "q5_0", "q5_1", "q5"),
+        ("q4_k_m", "q4_k_s", "q4_k", "q4_0", "q4_1", "q4"),
+        ("q3_k_m", "q3_k_s", "q3_k", "q3_0", "q3_1", "q3"),
+        ("q2_k", "q2_0", "q2_1", "q2"),
+        ("f32",),
+    ]
+    base = len(pref)
+    for gi, tokens in enumerate(groups):
+        for token in tokens:
+            if token in name:
+                return (base + gi, file_name)
+
+    return (base + len(groups), file_name)
 
 
 def main() -> None:

start.sh

@@ -11,6 +11,8 @@ export HF_HOME="${HF_HOME:-/data/hf-cache}"
 export MODEL_DIR="${MODEL_DIR:-/data/models}"
 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}"
 
 # Ensure data directories exist and are writable. Run as root in container startup.
 echo "preparing storage: $HF_HOME and $MODEL_DIR"
@@ -27,6 +29,30 @@ fi
 
 LLAMA_SERVER_BIN="${LLAMA_SERVER_BIN:-/usr/local/bin/llama-server}"
 
+# set server profile defaults
+case "$PERF_PROFILE" in
+  low_latency)
+    DEFAULT_BATCH=1
+    DEFAULT_UBATCH=1
+    DEFAULT_THREADS_BATCH=1
+    ;;
+  balanced)
+    DEFAULT_BATCH=8
+    DEFAULT_UBATCH=8
+    DEFAULT_THREADS_BATCH=${THREADS}
+    ;;
+  throughput)
+    DEFAULT_BATCH=512
+    DEFAULT_UBATCH=512
+    DEFAULT_THREADS_BATCH=${THREADS}
+    ;;
+  *)
+    DEFAULT_BATCH=8
+    DEFAULT_UBATCH=8
+    DEFAULT_THREADS_BATCH=${THREADS}
+    ;;
+esac
+
 server_args=(
   --model "$MODEL_PATH"
   --host 0.0.0.0
@@ -35,9 +61,9 @@ server_args=(
   --mmap
   -ngl 0
   -t "$THREADS"
-  --threads-batch "${THREADS_BATCH:-$THREADS}"
-  --batch-size "${BATCH_SIZE:-512}"
-  --ubatch-size "${UBATCH_SIZE:-512}"
+  --threads-batch "${THREADS_BATCH:-$DEFAULT_THREADS_BATCH}"
+  --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}"
@@ -49,5 +75,6 @@ if [[ -n "${PARALLEL:-}" ]]; then
 fi
 
 server_cmd=("$LLAMA_SERVER_BIN" "${server_args[@]}")
+echo "server command: $LLAMA_SERVER_BIN ${server_args[*]}"
 echo "starting llama-server"
 exec su -s /bin/sh appuser -c "$(printf '%q ' "${server_cmd[@]}")"