Support fixed ubatch batch sweep at 16K context

scripts/bench-ubatch-rocmi4.sh

@@ -1,9 +1,9 @@
 #!/usr/bin/env bash
-# bench-ubatch-rocmi4.sh — sweep ubatch size at 8k context on the
+# bench-ubatch-rocmi4.sh — sweep batch size with a fixed ubatch at 16k context
-# frank-default-large ROCmFPX standalone-MTP stack (dankclaw, gfx1151).
+# on the frank-default-large ROCmFPX standalone-MTP stack (dankclaw, gfx1151).
 #
-# Boots one llama-server instance per ubatch on a scratch port with a fixed
+# Boots one llama-server instance per batch on a scratch port with a fixed
-# -b and -c 32768, runs a standardized ~75%-of-ctx prose workload, and reads
+# ubatch and -c 16384, runs a standardized ~75%-of-ctx prose workload, and reads
 # authoritative prefill/generation tok/s from the server's own print_timing
 # lines. Steps from UBATCH_MIN to UBATCH_MAX inclusive by UBATCH_STEP.
 #
@@ -26,12 +26,9 @@ if [[ -f "$ENV_FILE" ]]; then
 fi
 
 # --- sweep parameters -------------------------------------------------------
-: "${UBATCH_MIN:=256}"
+: "${UBATCH_FIXED:=512}"
-: "${UBATCH_MAX:=6144}"
+: "${BATCH_VALUES:=512 1024 1536 2048}"
-: "${UBATCH_STEP:=256}"
+: "${CTX:=16384}"                # 16k context
-: "${UBATCH_VALUES:=256 512 768 1024 1280 1536 2048 3072 4096 5120 6144}"
-: "${CTX:=32768}"                # 32k contexts
-: "${BENCH_BATCH:=32768}"         # batch is fixed to the 32k context size
 : "${OUT_TOKENS:=256}"
 : "${RUNS:=3}"                   # decode runs per ubatch
 : "${PORT_BASE:=7894}"           # scratch ports 7894..
@@ -111,14 +108,15 @@ wait_health() {
     return 0
 }
 
-# run_one <ubatch> <port> <out_tokens> <runs> — boot one server instance, run
+# run_one <batch> <port> <out_tokens> <runs> — boot one server instance, run
 # the workload, parse print_timing, print a RESULT line.
 run_one() {
-    local ubatch="$1" port="$2" out_tokens="$3" runs="$4"
+    local batch="$1" port="$2" out_tokens="$3" runs="$4"
-    local log="/tmp/ubatch-${ubatch}-p${port}.log"
+    local ubatch="$UBATCH_FIXED"
+    local log="/tmp/ubatch-${ubatch}-batch-${batch}-p${port}.log"
     : > "$log"
 
-    echo "=== ubatch=${ubatch} ctx=${CTX} batch=${BENCH_BATCH} port=${port} runs=${runs} ==="
+    echo "=== ubatch=${ubatch} ctx=${CTX} batch=${batch} port=${port} runs=${runs} ==="
     "$LLAMA_BIN" \
         --model "$MODEL" \
         "${MMPROJ_ARGS[@]}" \
@@ -128,7 +126,7 @@ run_one() {
         --host 127.0.0.1 --port "$port" \
         -dev ROCm0 -ngl 999 -np 1 \
         -t 16 -tb 32 \
-        -c "$CTX" -b "$BENCH_BATCH" -ub "$ubatch" \
+        -c "$CTX" -b "$batch" -ub "$ubatch" \
         -fa on -ctk f16 -ctv f16 -ctxcp 16 \
         --cache-ram 8192 --cache-prompt --metrics --no-mmap \
         --spec-type draft-mtp --model-draft "$MTP_DRAFT_FILE" \
@@ -236,9 +234,9 @@ main() {
     stop_live_service
     local idx=0 ubatch port
     printf "%-10s %-14s %-10s\n" "UBATCH" "PREFILL" "GEN_TPS"
-    for ubatch in $UBATCH_VALUES; do
+    for batch in $BATCH_VALUES; do
         port=$((PORT_BASE + idx))
-        run_one "$ubatch" "$port" "$OUT_TOKENS" "$RUNS"
+        run_one "$batch" "$port" "$OUT_TOKENS" "$RUNS"
         idx=$((idx + 1))
     done
     echo "=== sweep complete ==="

tests/frank_large_config.bats

@@ -65,12 +65,10 @@ TEMPLATE_FILE="$REPO_ROOT/scripts/chat-template.jinja"
 
 @test "large ubatch bench defaults to the active model and requested sweep" {
     local bench="$REPO_ROOT/scripts/bench-ubatch-rocmi4.sh"
-    grep -q ': "${UBATCH_MIN:=256}"' "$bench"
+    grep -q ': "${UBATCH_FIXED:=512}"' "$bench"
-    grep -q ': "${UBATCH_MAX:=6144}"' "$bench"
+    grep -q 'BATCH_VALUES.*512 1024 1536 2048' "$bench"
-    grep -q 'UBATCH_VALUES.*256 512 768 1024 1280 1536 2048 3072 4096 5120 6144' "$bench"
     grep -q ': "${STOPSVC:=1}"' "$bench"
-    grep -q ': "${CTX:=32768}"' "$bench"
+    grep -q ': "${CTX:=16384}"' "$bench"
-    grep -q ': "${BENCH_BATCH:=32768}"' "$bench"
     grep -q 'MODEL_FILE' "$bench"
     grep -q 'REASONING_FORMAT' "$bench"
 }
@@ -82,6 +80,13 @@ TEMPLATE_FILE="$REPO_ROOT/scripts/chat-template.jinja"
     grep -qF 'eval time =\s' "$bench"
 }
 
+@test "large ubatch bench supports fixed ubatch batch sweep" {
+    local bench="$REPO_ROOT/scripts/bench-ubatch-rocmi4.sh"
+    grep -q ': "${UBATCH_FIXED:=512}"' "$bench"
+    grep -q 'BATCH_VALUES' "$bench"
+    grep -q -- '-b "$batch" -ub "$ubatch"' "$bench"
+}
+
 @test "frank-default-large uses the tracked Qwen-Sharp template on :7770" {
     grep -q '^CHAT_TEMPLATE_FILE=' "$ENV_FILE"
     grep -q -- '--jinja' "$SERVE_FILE"