From fb1b64908556c5b1d1a4a9e8cbe4bbcb06f098ca Mon Sep 17 00:00:00 2001 From: benoit lange Date: Mon, 5 Oct 2026 20:47:17 +0200 Subject: [PATCH] calibrate: with batch slots, measure the batch settings with every slot decoding at once calibrate.py tuned decode speed with one request at a time, which is the solo path: with "parallel": N (--batch) two more settings decide how fast the slots run together, and only a load of several requests at once shows them. Step 5 (only with batch slots) measures the total tok/s of every slot decoding at once for the config as it runs, without --batch-groups (with a layer split), and with half the --kv-resident cells (room for more experts in VRAM) - a restart each, two interleaved rounds - and keeps one only when it beats the config by more than MIN_GAIN, like the other steps. apply() writes those two as calibrated (None removes the flag); they have no product default to go back to. tools/test_calibrate.py: a stand-in engine with batch slots; the config is kept when flat, a smaller resident K/V and dropping the groups are found when faster, and no step 5 without slots. Co-Authored-By: Claude Opus 5.5 --- tools/calibrate.py | 72 +++++++++++++++++++++++++++++++++++++++++ tools/test_calibrate.py | 51 +++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+) diff --git a/tools/calibrate.py b/tools/calibrate.py index 7790056237e..7f85087129a 100644 --- a/tools/calibrate.py +++ b/tools/calibrate.py @@ -11,6 +11,12 @@ The first two are measured through one engine (per-request `strata_tune` keys); the worker count needs a restart per value. Decode speed only: the prompt path streams every expert whatever these settings say. +With batch slots ("parallel": N / --batch), two more settings decide how fast the slots run together, and only a +load of several requests at once shows them (one request takes the solo path): + --batch-groups with a layer split, the slots in groups pipelined through the GPUs (or all of them in one window). + --kv-resident the K/V cells kept in VRAM per sequence; fewer leave room for more experts in the cache. +They are measured last, a restart each, as the total tok/s of every slot decoding at once. + A setting is kept only when it beats the default by more than MIN_GAIN in an interleaved re-measurement - the adaptive expert tier and the OS make single measurements noisy by a few percent. @@ -156,6 +162,7 @@ def measure(base_args: list[str], ids_list, start_engine, say=print) -> dict: d_pcie = float(info.get("pcie_frac", 0.55)) d_minp = float(info.get("spec_min_p", 0.5)) d_workers = int(info.get("pool_workers", 0)) or None + slots = int(getattr(eng, "batch", 0) or info.get("batch_slots", 0) or 0) s = Session(eng, ids_list) s.warm_up() # 1. the PCIe share, at the default draft floor @@ -210,11 +217,72 @@ def measure(base_args: list[str], ids_list, start_engine, say=print) -> dict: base_rate = statistics.median(by_workers[w_best]) elif by_workers.get(d_workers): base_rate = statistics.median(by_workers[d_workers]) + # 5. with batch slots: the batch settings, with every slot decoding at once (a restart each) + if slots > 1: + tuned = apply(base_args, settings) + p_settings, p_report = measure_parallel(tuned, slots, ids_list, start_engine, say) + settings.update(p_settings) + report.update(p_report) report["seconds"] = round(time.time() - t0) report["tok_s"] = round(base_rate, 1) if base_rate else None return {"settings": settings, "report": report} +def parallel_rate(engine, ids_list, clients: int) -> float: + """Total decode tok/s with `clients` requests at once: the engine runs them together in its batch slots.""" + counts = [0] * clients + + def one(i): + ids = ids_list[i % len(ids_list)] + counts[i] = sum(1 for t in engine.generate(ids, MAX_NEW, {"temperature": 0}, threading.Event()) if t is not None) + + t0 = time.time() + th = [threading.Thread(target=one, args=(i,)) for i in range(clients)] + for x in th: + x.start() + for x in th: + x.join() + return sum(counts) / max(time.time() - t0, 1e-9) + + +def parallel_candidates(args: list[str]) -> dict: + """The batch settings worth a restart each, from the config as it runs: name -> args.""" + out = {"as configured": list(args)} + groups = arg_value(args, "--batch-groups") + if groups and int(groups) > 1: + out["no --batch-groups"] = with_arg(args, "--batch-groups", None) + kvres = arg_value(args, "--kv-resident") + if kvres and int(kvres) > 1: + out[f"--kv-resident {int(kvres) // 2}"] = with_arg(args, "--kv-resident", str(int(kvres) // 2)) + return out + + +def measure_parallel(tuned_args: list[str], slots: int, ids_list, start_engine, say=print) -> tuple[dict, dict]: + """Step 5: every batch setting (a restart each) with all the slots decoding at once; -> (settings, report).""" + cands = parallel_candidates(tuned_args) + if len(cands) < 2: + return {}, {} + by = {k: [] for k in cands} + for _ in range(2): # interleaved: two rounds over every candidate + for name, args in cands.items(): + say(f" Measuring {slots} requests at once, {name} (restarts the engine) ...") + e = start_engine(args) + try: + parallel_rate(e, ids_list, slots) # warm-up: the slots' first windows and the expert tier + by[name].append(parallel_rate(e, ids_list, slots)) + finally: + close(e) + say(f" {slots} at once, {name}: {by[name][-1]:.1f} tok/s") + best = pick(by, "as configured") + settings = {} + if best != "as configured": + a = cands[best] + for flag in ("--batch-groups", "--kv-resident"): + if arg_value(a, flag) != arg_value(tuned_args, flag): + settings[flag] = arg_value(a, flag) + return settings, {"slots": slots, "parallel": {k: [round(x, 1) for x in v] for k, v in by.items()}} + + def close(eng): proc = getattr(eng, "proc", None) if proc is None: @@ -229,6 +297,7 @@ def close(eng): DEFAULTS = {"--pcie-frac": None, "--spec-min-p": "0.5", "--pool-workers": None} # None: the engine's own choice +BATCH_FLAGS = ("--batch-groups", "--kv-resident") # step 5: changed only when a calibration measured them better def apply(args: list[str], settings: dict) -> list[str]: @@ -238,6 +307,9 @@ def apply(args: list[str], settings: dict) -> list[str]: out = list(args) for flag, default in DEFAULTS.items(): out = with_arg(out, flag, settings.get(flag, default)) + for flag in BATCH_FLAGS: # the batch settings: as calibrated (None: without the flag) + if flag in settings: + out = with_arg(out, flag, settings[flag]) return out diff --git a/tools/test_calibrate.py b/tools/test_calibrate.py index a03286c5ad1..9c46461506b 100644 --- a/tools/test_calibrate.py +++ b/tools/test_calibrate.py @@ -42,6 +42,26 @@ def generate(self, ids, max_new, sampling, cancel): BASE = ["--pack", "p", "--spec", "4", "--spec-min-p", "0.5", "--max-context", "8192"] +BATCH = BASE + ["--batch", "8", "--batch-groups", "4", "--kv-resident", "32768"] + + +class FakeBatchEngine(FakeEngine): + """With batch slots: requests at once share the throughput `batch_speed(args)` (tok/s for all of them).""" + + def __init__(self, args, speed, batch_speed, starts=None): + super().__init__(args, speed, 6, starts) + self.batch = int(CAL.arg_value(args, "--batch") or 0) + self.batch_speed = batch_speed(args) + self.lock = __import__("threading").Lock() + + def generate(self, ids, max_new, sampling, cancel): + if sampling.get("strata_tune"): + yield from super().generate(ids, max_new, sampling, cancel) + return + import time as _t + _t.sleep(max_new * self.batch / self.batch_speed) # each of the `batch` requests gets batch_speed / batch + for _ in range(max_new): + yield 1 class Calibrate(unittest.TestCase): @@ -82,6 +102,37 @@ def test_old_calibration_is_the_baseline_reset(self): self.assertIsNone(CAL.arg_value(starts[0], "--pool-workers")) self.assertEqual(CAL.arg_value(starts[0], "--spec-min-p"), "0.5") + def run_batch(self, batch_speed): + starts = [] + with mock.patch.object(CAL, "MAX_NEW", 40): + res = CAL.measure(BATCH, [[1, 2, 3]] * 3, lambda a: FakeBatchEngine(a, lambda f, p, w: 50.0, batch_speed, starts), + say=lambda *_: None) + return res, starts + + def test_parallel_step_keeps_the_config_when_flat(self): + res, starts = self.run_batch(lambda a: 400.0) + self.assertEqual(res["settings"], {}) + self.assertEqual(res["report"]["slots"], 8) + self.assertEqual(set(res["report"]["parallel"]), {"as configured", "no --batch-groups", "--kv-resident 16384"}) + # the sweep, one start per worker count, then two interleaved rounds of the three batch candidates + self.assertEqual(len(starts), 1 + len(CAL.worker_candidates(6)) + 2 * 3) + + def test_parallel_step_finds_a_smaller_resident_kv(self): + res, _ = self.run_batch(lambda a: 500.0 if CAL.arg_value(a, "--kv-resident") == "16384" else 400.0) + self.assertEqual(res["settings"], {"--kv-resident": "16384"}) + a = CAL.apply(BATCH, res["settings"]) + self.assertEqual(CAL.arg_value(a, "--kv-resident"), "16384") + self.assertEqual(CAL.arg_value(a, "--batch-groups"), "4") + + def test_parallel_step_can_drop_the_groups(self): + res, _ = self.run_batch(lambda a: 500.0 if "--batch-groups" not in a else 400.0) + self.assertEqual(res["settings"], {"--batch-groups": None}) + self.assertNotIn("--batch-groups", CAL.apply(BATCH, res["settings"])) + + def test_no_parallel_step_without_slots(self): + res, starts = self.run_with(lambda f, p, w: 50.0) + self.assertNotIn("parallel", res["report"]) + def test_apply(self): a = CAL.apply(BASE, {"--pcie-frac": "0.35", "--pool-workers": "4"}) self.assertEqual(CAL.arg_value(a, "--pcie-frac"), "0.35")