Upload k3_bootstrap.sh with huggingface_hub

k3_bootstrap.sh

@@ -356,7 +356,26 @@ log "gpus=$NGPU, letting common_fit_params place everything"
 # cudaMalloc OOM, device 0 untouched. With autofit, both cards fill to 95%.
 COMMON=(
   -m "$FIRST"
-  --ctx-size 16384
+  # Long context is unusually cheap here, and the config says why:
+  #   kv_lora_rank 512 + qk_rope_head_dim 64 = 576 values per token per layer
+  #   full_attn_layers = every 4th layer   -> only 24 of 93 layers cache at all
+  #   the other 69 are KDA: recurrent state, constant regardless of context
+  # So the MLA cache costs 24 * 576 * 2 = 27.6 KB per token:
+  #   65536 ctx -> 1.8 GB    262144 -> 7.2 GB    1048576 -> 29 GB
+  # against ~38 GB of VRAM left free by the weights on a 5-GPU host. 16384 was
+  # never a hardware limit, just an untested default -- and it could not even
+  # accept one Claude Code request, whose system prompt alone exceeds 32k.
+  # What actually constrains context is the weights sharing the same cards, so
+  # check the VRAM dump printed after READY: if the five stop being ~38-40 GB
+  # full, autofit has started spilling layers to the CPU and decode collapses.
+  --ctx-size "${K3_CTX:-262144}"
+  # Prefill is compute-bound and batches, unlike decode. The default micro
+  # batch of 512 leaves the GPUs under-fed on the long prompts this client
+  # sends; 2048 gives the MoE GEMMs four times the rows to amortise each
+  # expert weight read over. It costs compute-buffer VRAM, which is the other
+  # thing to watch in that same dump.
+  --batch-size "${K3_BATCH:-4096}"
+  --ubatch-size "${K3_UBATCH:-2048}"
   --flash-attn on
   --jinja
   # Not carried by the GGUF metadata; the model mis-routes experts without them.
@@ -374,12 +393,17 @@ COMMON=(
   # bottleneck: if it were, halving it would have nearly doubled throughput.
   --override-kv "kimi-k3.expert_used_count=int:${K3_EXPERTS:-16}"
   --host 0.0.0.0 --port "$PORT"
-  # Batch-1 decode is the worst case for a MoE model: 16 experts are read from
-  # VRAM to produce a single token, so the weight traffic is amortised over one
-  # row. Serving several streams from one forward pass reuses that traffic.
-  # Slots share --ctx-size (16384/4 = 4096 each), so this costs no extra KV
-  # VRAM. The CONCURRENCY phase below measures whether it actually pays.
-  --parallel "${K3_PARALLEL:-4}"
+  # Measured on this pod, 64 tokens per stream, same server:
+  #   1 stream    5.56 tok/s aggregate
+  #   2 streams   9.85 tok/s aggregate  (1.77x)
+  #   4 streams  14.53 tok/s aggregate  (2.61x)
+  # Concurrency scales, so decode is bound by serialisation, not bandwidth --
+  # at batch 1 the effective read rate is ~51 GB/s against 696 GB/s available
+  # on the active card. There is idle silicon and extra streams harvest it.
+  # But slots SPLIT --ctx-size, and a Claude Code client needs the whole
+  # context for one stream, so the default is 1. Raise K3_PARALLEL (and
+  # K3_CTX with it) only when serving several clients at once.
+  --parallel "${K3_PARALLEL:-1}"
   --threads "$THREADS"
   --threads-batch "$THREADS"
   # Greedy decoding measurably collapses this ~1-bit quant into repetition
@@ -638,49 +662,40 @@ curl -s -m 900 "http://127.0.0.1:$PORT/v1/chat/completions" -H 'Content-Type: ap
   | jq -r 'if .error then "ERROR: \(.error.message)" else (.choices[0].message.content // "EMPTY CONTENT"), (.choices[0].message.reasoning_content // empty) end'
 echo
 
-# ---------------------------------------------------------------- concurrency
-# The one question the placement experiments never asked. Six placements moved
-# 60-88 GB of weights between CPU and GPU and barely changed single-stream
-# throughput; the host vCPU count changed it 1.9x and full VRAM residency
-# another 1.4x. What none of them tested is whether the hardware is *idle*
-# between tokens.
+# ------------------------------------------------------------------- prefill
+# The number that decides whether this serves a Claude Code client. Its system
+# prompt alone is 32k+ tokens, and every selftest prompt above is 6-22 tokens,
+# so their "prefill tok/s" is pure per-request overhead and says nothing.
 #
-# The discriminator is aggregate throughput vs concurrency:
-#   scales ~linearly  -> decode is latency/serialisation bound (llama.cpp splits
-#                        layers across GPUs, so at batch 1 exactly one GPU
-#                        computes while the others wait). Idle silicon exists
-#                        and concurrency harvests it.
-#   stays flat        -> genuinely bandwidth bound; nothing left on the table
-#                        short of a different quant or fewer GPUs' worth of
-#                        weights.
-# Either answer is worth having, and it costs one server start.
-echo "===== CONCURRENCY ====="
-CONC_PROMPT='Write a detailed explanation of how Mixture-of-Experts routing works:'
-CONC_N=64
-for C in 1 2 4; do
-  [ "$C" -gt "${K3_PARALLEL:-4}" ] && continue
-  t0=$(date +%s.%N)
-  for i in $(seq 1 "$C"); do
-    # Distinct prefix per stream: identical prompts would share the prefix
-    # cache and make prefill look free, inflating the aggregate.
-    jq -n --arg p "[$i] $CONC_PROMPT" --argjson n "$CONC_N" \
-      '{prompt:$p,n_predict:$n,temperature:0,repeat_penalty:1.0,stream:false}' \
-      | curl -s -m 900 "http://127.0.0.1:$PORT/completion" \
-          -H 'Content-Type: application/json' -d @- >"/tmp/conc.$i.json" &
-  done
-  wait
-  t1=$(date +%s.%N)
-  gen=0
-  for i in $(seq 1 "$C"); do
-    n=$(jq -r '(.timings // {}).predicted_n // 0' "/tmp/conc.$i.json" 2>/dev/null)
-    gen=$((gen + ${n:-0}))
-  done
-  awk -v c="$C" -v g="$gen" -v a="$t0" -v b="$t1" \
-    'BEGIN{d=b-a; printf "concurrency %d: %d tok in %.1f s = %.2f tok/s aggregate (%.2f per stream)\n", c, g, d, g/d, g/d/c}'
-done
+# Two costs matter and they are three orders of magnitude apart:
+#   cold  - the first turn, which must actually run the whole prompt
+#   warm  - every turn after it, which llama.cpp serves from the slot's prompt
+#           cache and re-runs only the tokens that changed
+# Measured on 5x A40, 3408-token prompt: 167 tok/s cold, and the warm repeat
+# reprocessed 4 tokens in 0.53 s. Prefill batches (unlike decode), so it is
+# ~23x the decode rate; a 32k system prompt costs ~3 min once, then ~nothing.
+echo "===== PREFILL ====="
+PF_JSON=/tmp/pf.json
+python3 - "$PF_JSON" "${K3_PF_TOKENS:-32000}" <<'PYEOF'
+import json, sys
+out, target = sys.argv[1], int(sys.argv[2])
+unit = ("Section %d. The routing network assigns each token to a sparse subset "
+        "of experts and the gate weights are normalised over the selected set, "
+        "which keeps activation cost constant. ")
+# ~34 tokens per unit as tokenised by this model (measured: 100 units = 3408).
+s = "".join(unit % i for i in range(1, max(1, target // 34) + 1))
+s += "
+Summarise the above in one sentence:"
+json.dump({"prompt": s, "n_predict": 4, "temperature": 0, "stream": False}, open(out, "w"))
+PYEOF
+pf_run() {
+  curl -s -m 3600 "http://127.0.0.1:$PORT/completion" -H 'Content-Type: application/json'        --data-binary @"$PF_JSON"     | jq -r '(.timings // {}) | "  \(.prompt_n // 0) tok in \(((.prompt_ms // 0)/1000)*10|round/10) s = \((.prompt_per_second // 0)*10|round/10) tok/s"'
+}
+echo "cold (prompt never seen):"; pf_run
+echo "warm (identical prompt, should hit the slot prompt cache):"; pf_run
 echo
 
-log "RESULT selftest done - mode=$MODE parallel=${K3_PARALLEL:-4}"
+log "RESULT selftest done - mode=$MODE parallel=${K3_PARALLEL:-1} ctx=${K3_CTX:-262144}"
 log "server live on port $PORT (/completion and /v1/chat/completions)"
 
 # Keep serving, and hot-reload on script changes WITHOUT a pod restart.