From e09a1fe0ab046b44619e7cffd1fceae6a56f8f8c Mon Sep 17 00:00:00 2001 From: rasa04 Date: Sun, 25 Oct 2026 03:32:45 +0500 Subject: [PATCH] [Engine] Mac: calibrate the decode kernels' tiles on the device (setup.sh --calibrate) The Mac's part of tools/calibrate.py. The Metal decode kernels (the dense GEMVs, gr_down, the routed experts' gate/up and down) run on tiles tuned on the M4 and M5; a GPU with less ALU per byte of bandwidth wants others (M1 Max: a 2-token window's dense GEMVs 33 -> 21 ms on per-shape best tiles). - Engine: STRATA_METAL_CALIBRATE= times every candidate tile of each GEMV shape and gr_down kind when a window size is first encoded, and the experts' tiles (per pair of formats) after that window's first run; outputs go to a scratch buffer, so the model's state is not touched. A tile replaces the default only when it measured >= 3% faster. STRATA_METAL_TILES= reads them and never measures; a shape the file lacks, or a file from another GPU (its first line names the device), keeps the defaults. Neither set: unchanged. - tools/calibrate.py on a Mac: the measuring start, then the tiles race the defaults on the config's own arguments (a restart each, interleaved rounds) and go into the config's env only when they win by more than MIN_GAIN. One file a model. metal/setup_mac.py --calibrate applies it; a reinstall keeps it. - The fp16 decode (STRATA_METAL_F16=1) keeps its tiles: the calibration neither measures its GEMVs and routed experts nor gives them a tile measured on the fp32 kernels. - setup.sh --calibrate does not start beside a running engine; strata engine calibrate refuses a Mac config (its tiles would replace the file the config names). --- .gitignore | 1 + docs/MAC.md | 10 +- metal/setup_mac.py | 41 +++- metal/src/engine/window.cpp | 369 ++++++++++++++++++++++++++++- strata_control/calibration.py | 2 + tests/terminal/test_calibration.py | 8 + tools/calibrate.py | 60 ++++- tools/test_calibrate.py | 56 ++++- 8 files changed, 536 insertions(+), 11 deletions(-) diff --git a/.gitignore b/.gitignore index 65457fdc042..7d2f3c4f7bc 100644 --- a/.gitignore +++ b/.gitignore @@ -36,6 +36,7 @@ td.log /run-*.sh /strata-*.json /strata-*.log +/metal-tiles*.txt /build-strata.bat /build-strata-cuda12.bat /build-cuda12/ diff --git a/docs/MAC.md b/docs/MAC.md index e7342ade9e4..37c65f8a03c 100644 --- a/docs/MAC.md +++ b/docs/MAC.md @@ -226,7 +226,15 @@ from each other: top-1 agreement 95.4% against 95.8%, mean KL 0.0207 against 0.0 - **Context:** measured up to 32K (needles, KL). Setup recommends 32K below 96 GB and 64K from 96 GB. Up to the model's 262K can be set, but it is not measured. - **Several requests at once.** `--parallel N` keeps N sequences side by side in the engine, each with its own running state and K/V (0.65 GB each at a 36,864-token context). Measured on an M5 Max with IQ3_S, the MTP draft layer and `--spec 2`, N requests of 150 tokens sent at the same time, tokens a second over all of them: 1 request 85, 2 requests 95 to 102, 3 requests 115, 4 requests 119 to 125. One request alone decodes at 96 to 100; with 2, 3 and 4 at once each decodes at 53 to 61, 41 to 48 and 32 to 40. The engine uses 8.3, 12.0, 14.0 and 16.0 GB for 1 to 4. -- **Not ported:** the control vector (`--control-vector-scaled`) and the PC tuning (`--calibrate`). +- **Not ported:** the control vector (`--control-vector-scaled`). +- **`./setup.sh --calibrate`:** the engine measures the tiles of its decode kernels on this Mac (the dense GEMVs, the + hyper-connection read, and the routed experts when the model is in memory - not when they stream from the SSD; the + defaults were tuned on the M4 and M5), then they race the defaults and go into the config's `env` (`STRATA_METAL_TILES`) + only when they are over 3% faster. The fp16 kernels (`STRATA_METAL_F16=1`) keep their default tiles. The server must be stopped. + A few minutes with the model in memory, longer with streaming; again after an update that changes the kernels. The + answers may differ slightly from the defaults' (the order of a sum: a near-tie can flip). Measured on an M1 Max only, + with fp16 off: decode +30-32% on short chats, +26% after a 6K-token prompt; with fp16 on: +6.8% and +2.7% in two + runs (the second under the 3% bar: the defaults stayed). - **Streaming** is measured at 32 GB only. - **Families:** the original Qwen3.8-Flash-Next, Unsloth's UD-Q4/Q5/Q6_K_XL and agentionai's Gyro-S are measured. Setup repacks Unsloth's shards into one file for the GPU; a Gyro file is used as it is. The Swift and Coder releases are not set up on Mac diff --git a/metal/setup_mac.py b/metal/setup_mac.py index 2804d93214a..f55440cc21f 100644 --- a/metal/setup_mac.py +++ b/metal/setup_mac.py @@ -330,6 +330,33 @@ def start(cfg_path: Path, port: int | None, open_browser: bool = True) -> int: # ------------------------------------------------------------------------------------------------ main +def calibrate_config(cfg_path: Path) -> bool: + """The Mac's tuning (tools/calibrate.py): the engine measures its GPU kernels' tiles, they race the defaults, and the + config's "env" names them (STRATA_METAL_TILES) only when they are faster.""" + running = engine_pids() + if running: # a second engine beside it can freeze this Mac (repack_running) + S.warn(f"a Strata engine is running (pid {running[0]}): stop the server (close its window), then run " + "./setup.sh --calibrate; the config is unchanged") + return False + sys.path.insert(0, str(ROOT / "tools")) + import calibrate as CAL + cfg = json.loads(cfg_path.read_text(encoding="utf-8-sig")) + S.say() + S.say(" Tuning Strata for this Mac: the engine measures its GPU kernels' tiles, then they race the defaults.") + S.say(" It takes a few minutes (longer when the experts stream from the SSD); the Mac is busy meanwhile.") + try: + res = CAL.run(cfg, say=S.say) + except Exception as e: # never stops an install: the defaults stay + S.warn(f"the tuning did not finish ({e}): the config is unchanged") + return False + cfg["env"] = CAL.apply_env(cfg.get("env"), res["env"]) + S.write_config(cfg_path, cfg) + tok_s = res["report"].get("tok_s") + S.ok(("tuned for this Mac: the measured tiles" if res["env"] else "tuned for this Mac: the default tiles are already the fastest here") + + (f" ({tok_s} tok/s)" if tok_s else "")) + return True + + def main(argv=None) -> int: ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) ap.add_argument("--family", choices=["qwen", "unsloth", "gyro"], default=None, @@ -352,6 +379,8 @@ def main(argv=None) -> int: help="where model files come from: auto (huggingface), huggingface, modelscope; " "same as STRATA_SOURCE") ap.add_argument("--setup", action="store_true") + ap.add_argument("--calibrate", action="store_true", + help="measure the GPU kernels' tiles on this Mac (a few minutes) and keep them when they are faster") ap.add_argument("--no-start", action="store_true") ap.add_argument("--check", action="store_true") ap.add_argument("--update", action="store_true") @@ -373,6 +402,8 @@ def main(argv=None) -> int: models_root = Path(a.models_dir) if a.models_dir else data / "models" have = mac_configs() + if have and a.calibrate and not (a.setup or a.model or a.check or a.update or a.build): + return 0 if calibrate_config(have[0]) else 1 if have and not (a.setup or a.model or a.check or a.no_start or a.update or a.build): return start(have[0], a.port, a.browser is not False) @@ -553,6 +584,13 @@ def main(argv=None) -> int: cfg = mac_config(exe, engine_args(pack, shards, rt, ctx, stream, info["ram_gb"], model), pack, f"{fam['name']}-{model.lower()}", tag, port) cfg["env"] = {"STRATA_METAL_STREAM": "1" if stream else "0", # the engine follows setup's measured choice **_fit.engine_env(model, stream)} # STRATA_METAL_F16=1 for a resident IQ3_S / Q2_0 (mac_fit.F16_MODELS) + cfg_path = ROOT / f"strata-{tag.lower()}.json" + try: # an earlier --calibrate survives a reinstall + tiles = (json.loads(cfg_path.read_text(encoding="utf-8-sig")).get("env") or {}).get("STRATA_METAL_TILES") + except (OSError, ValueError): + tiles = None + if tiles and Path(tiles).exists(): + cfg["env"]["STRATA_METAL_TILES"] = tiles if vision == "yes": cfg["args"] += ["--vision"] cfg["vision"] = {"exe": str(vexe), "mmproj": str(mmproj), "model": str(files[0]), "gpu": vision_gpu, @@ -565,8 +603,9 @@ def main(argv=None) -> int: cfg["draft_vocab"] = a.draft_vocab if a.browser is not None: cfg["open_browser"] = a.browser - cfg_path = ROOT / f"strata-{tag.lower()}.json" S.write_config(cfg_path, cfg) + if a.calibrate: + calibrate_config(cfg_path) script = S.write_run_script(tag, cfg_path, port, cfg.get("open_browser") is not False) S.ok(f"start script: {script.name}") S.say() diff --git a/metal/src/engine/window.cpp b/metal/src/engine/window.cpp index 6f40b6c4d78..a4494a05434 100644 --- a/metal/src/engine/window.cpp +++ b/metal/src/engine/window.cpp @@ -31,6 +31,7 @@ namespace strata::gpu { void* host_view(const void* device_pointer); } #include "strata/kernels/cpu/expert_layout.hpp" #include +#include #include #include #include @@ -141,6 +142,59 @@ bool gemv16_tile(int L, int R, int S, int K) { // the writers of the window's x: their pipelines with the half planes (function constant 25) under STRATA_METAL_F16 mtl::Constants x16c(mtl::Constants c = {}) { return metal_f16() ? c.b(25, true) : c; } std::string x16k(const std::string& k) { return metal_f16() ? k + "#h" : k; } +// The decode kernels' tiles measured on this GPU - the Mac's part of the calibration (tools/calibrate.py): the dense +// GEMVs, gr_down and the routed experts' gate/up and down. STRATA_METAL_CALIBRATE=: measure them while the engine +// answers (Window::calibrate), into . STRATA_METAL_TILES= (what the calibration puts in the config's "env"): +// use them; a shape the file lacks, or a file from another GPU, keeps the defaults. The defaults are the M4's and M5's, +// and a GPU with less ALU per byte of bandwidth (the M1 Max: 33 -> 21 ms of a 2-token window's GEMVs) wants other tiles. +// A tile replaces the default only when it measured >= 3% faster; the answers may differ slightly (the order of a sum). +const char* metal_env_file(const char* name) { + const char* e = std::getenv(name); + return e && *e ? e : nullptr; +} +const char* metal_calibrate_file() { + static const char* f = metal_env_file("STRATA_METAL_CALIBRATE"); + return f; +} +const char* metal_tiles_file() { + static const char* f = metal_calibrate_file() ? metal_calibrate_file() : metal_env_file("STRATA_METAL_TILES"); + return f; +} +// the candidates: every tile gemv.metal instantiates (GV_SET), and the formats it instantiates them for +const std::array kGemvTiles[] = {{32, 1, 4, 1}, {32, 2, 4, 1}, {16, 1, 4, 1}, {16, 2, 4, 1}, {8, 1, 4, 1}, + {8, 2, 4, 1}, {4, 2, 4, 1}, {4, 1, 4, 1}, {32, 1, 4, 4}, {32, 2, 4, 2}, + {32, 1, 8, 8}, {16, 1, 8, 4}, {32, 1, 8, 2}, {8, 1, 8, 1}, {1, 1, 4, 1}, + {1, 2, 4, 1}, {2, 1, 4, 1}, {2, 2, 4, 1}, {1, 1, 2, 1}, {4, 4, 4, 1}}; +bool gemv_tunable(int ty) { + switch (ty) { + case ref::F32: case ref::BF16: case ref::Q8_0: case ref::Q4_0: case ref::Q5_0: case ref::IQ4_NL: case ref::IQ4_XS: + case ref::Q2_0: case 142 /* Q2_0R */: case ref::IQ3_S: case ref::Q3_K: case ref::Q4_K: case ref::Q5_K: case ref::Q6_K: + return true; + default: return false; + } +} +// a tile fits a row of n_in when its K split divides the units and each lane of a row gets at least one unit pair +bool gemv_tile_ok(const std::array& c, int ty, int n_in) { + const int unit = ty == ref::Q2_0 || ty == 142 || ty == ref::IQ3_S || ty == ref::BF16 || ty == ref::F32 ? 16 : 32; + const int U = n_in / unit; + return U % c[3] == 0 && U / c[3] >= c[0] / 2; +} +// gr_down's variants (gr.metal GRD): simdgroups a threadgroup, the K split of a row +const std::array kGrdVariants[] = {{8, 8}, {8, 4}, {4, 4}, {4, 2}, {16, 16}, {16, 8}}; +// the routed experts' tiles (moe.metal GU_SET / DN_SET: lanes a row, rows a lane, simdgroups) and their formats +const std::array kGuTiles[] = {{32, 1, 4}, {32, 2, 4}, {32, 4, 4}, {32, 2, 2}, {32, 2, 8}, {32, 1, 8}, + {16, 1, 4}, {16, 2, 4}, {8, 1, 4}, {8, 2, 4}, {8, 1, 2}}; +const std::array kDnTiles[] = {{8, 1, 4}, {8, 2, 4}, {8, 4, 4}, {4, 1, 4}, {4, 2, 4}, {4, 4, 4}, + {8, 1, 8}, {8, 2, 8}, {16, 1, 4}, {16, 2, 4}, {32, 1, 4}, {2, 2, 4}}; +bool moe_tunable(bool down, int ty) { + if (down) return ty == ref::Q5_1 || ty == ref::Q5_0 || ty == ref::Q8_0 || ty == ref::Q2_0 || ty == 142 || ty == ref::IQ4_NL; + switch (ty) { + case ref::Q2_0: case ref::IQ3_XXS: case ref::IQ2_XXS: case ref::IQ2_XS: case ref::IQ2_S: case ref::IQ3_S: + case ref::IQ4_XS: case ref::Q4_K: case ref::Q5_K: case ref::Q6_K: case ref::Q8_0: + return true; + default: return false; + } +} int xkind(int t) { switch (t) { @@ -458,9 +512,26 @@ class Window { std::vector ds; std::vector cuts; // SSD streaming: the segment ends (one per layer, before its experts) std::vector> qsa; // (dispatch, kind): 0 QsaArgs n_tok = T, 1 the merge (MAXT), 2 the scores + // STRATA_METAL_CALIBRATE: the dispatches whose tile is not measured yet (calibrate) + struct Gv { int i; std::array shape, tile; }; // (type, n_in, n_out, ntok), the default (L, R, S, K) + struct Grd { int i; std::string pre; int ntok; }; + std::vector gv; + std::vector grd; + // a layer's routed experts: the route dispatch (and the index of its MoeArgs), gate/up, down; their formats and tiles + struct Moe { int r, ra, g, d, gt, dt, ntok; bool ct; std::array gu, dn; }; + std::vector moe; }; std::map wcache_; bool f16_enc_ = false; // STRATA_METAL_F16: a decode window is being encoded (the MTP layer stays fp32) + Cached* cal_rec_ = nullptr; // the encode in progress (calibrate) + std::map, std::array> cal_gemv_; // (type, n_in, n_out, ntok) -> (L, R, S, K) + std::map, int> cal_grd_; // (gr_down kind, ntok) -> kGrdVariants index + std::map, std::array> cal_moe_; // (0 gate/up | 1 down, type, ntok) -> (L, R, S) + void calibrate_moe(Cached& c); + void calibrate_save(); + std::string cal_file_; + void calibrate(Cached& c); + void calibrate_load(); std::vector>* rec_ = nullptr; // ---- SSD streaming (stream.h): the dense weights in W_, the routed experts in the arena std::unique_ptr st_; @@ -693,6 +764,7 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta lib_ = cx_->compile({"quant.metal", "units.metal", "moe.metal", "gemv.metal", "router.metal", "gr.metal", "gr2.metal", "glue.metal"}, metal_f16() ? "#define STRATA_METAL_F16 1\n" : ""); rcounter_ = cx_->alloc(16); + if (metal_tiles_file()) calibrate_load(); lgdn_ = cx_->compile({"gdn.metal"}, "#define NO_GGML_COMMON 1\n"); lqsa_ = cx_->compile({"qsa.metal", "qsa2.metal"}, "#define NO_GGML_COMMON 1\n"); lsmp_ = cx_->compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n"); @@ -901,6 +973,11 @@ void Window::gemv(std::vector& ds, const gguf::Tensor& t, const A } } const int k16 = f16_enc_ && x.n == n_in ? x16_kind(ty) : -1; + if (metal_tiles_file() && gemv_tunable(ty) && k16 < 0) { // the fp16 GEMVs keep their tiles (not measured) + const auto it = cal_gemv_.find({ty, n_in, n_out, ntok}); + if (it != cal_gemv_.end()) { L = it->second[0]; Rr = it->second[1]; S = it->second[2]; K = it->second[3]; } + else if (cal_rec_) cal_rec_->gv.push_back({(int) ds.size(), {ty, n_in, n_out, ntok}, {L, Rr, S, K}}); + } const bool f16 = k16 >= 0 && gemv16_tile(L, Rr, S, K); const std::string fn = std::string(f16 ? "gemv16_" : "gemv_") + std::to_string(ty) + "_" + std::to_string(L) + "_" + std::to_string(Rr) + "_" + std::to_string(S) + "_" + std::to_string(K); @@ -957,12 +1034,25 @@ void Window::gr_read(std::vector& ds, const std::string& pre, boo const bool q8 = is_q8(tdn) && is_q8(tup); const bool injf = q8 && !final_ && !q8off_.count(tin.name) && tin.type == ref::F32; mtl::Dispatch d2; - const char* dn = injf ? "gr_down_q8f_8_8" : q8 ? "gr_down_q8_8_8" : "gr_down_8_8"; - d2.p = P(lib_, dn, mtl::Constants().i(1, ntok), std::string(dn) + "#" + std::to_string(ntok)); + // the variant: 8_8 (8 simdgroups split a row), or the one measured fastest here (STRATA_METAL_TILES) + const std::string dpre = injf ? "gr_down_q8f_" : q8 ? "gr_down_q8_" : "gr_down_"; + int var = 0; + if (metal_tiles_file()) { + const auto it = cal_grd_.find({dpre, ntok}); + if (it != cal_grd_.end()) { + const auto [sg, ks] = kGrdVariants[it->second]; + const std::string fn = dpre + std::to_string(sg) + "_" + std::to_string(ks); + if (P(lib_, fn, mtl::Constants().i(1, ntok), fn + "#" + std::to_string(ntok))->max_threads >= 32 * sg) + var = it->second; // (a variant above this pipeline's threadgroup limit keeps the default: require_fit) + } else if (cal_rec_) cal_rec_->grd.push_back({(int) ds.size(), dpre, ntok}); + } + const auto [grd_sg, grd_ks] = kGrdVariants[var]; + const std::string dn = dpre + std::to_string(grd_sg) + "_" + std::to_string(grd_ks); + d2.p = P(lib_, dn, mtl::Constants().i(1, ntok), dn + "#" + std::to_string(ntok)); d2.args = {mtl::Arg::B(*bd, od), mtl::Arg::B(*bi, oi), mtl::Arg::B(xn_), mtl::Arg::B(lo_), mtl::Arg::B(inj_[1 - cur_inj_]), mtl::Arg::V(a)}; - d2.grid[0] = LR + HC; - d2.tg[0] = 256; + d2.grid[0] = (LR + HC + grd_sg / grd_ks - 1) / (grd_sg / grd_ks); // RS = SG / KS rows a threadgroup + d2.tg[0] = 32 * grd_sg; push(ds, d2); mtl::Dispatch d3; const char* up = q8 ? "gr_up_q8" : "gr_up"; @@ -1189,13 +1279,26 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 // the trellis gate/up: 16 lanes a row, 2 rows a lane (expert_parity --bench --sweep, M4, TQK6, 10 experts: 166.6 us a layer // against 174.5 for 32,2,4 and 196.8 for 32,1,4; also the fastest at 2 and 3 tokens) const bool tq_ = tg_.type >= 144; - const int gl = tq_ ? 16 : 32, gr = tq_ ? 2 : 1, gs = 4, ds_ = 4; + int gl = tq_ ? 16 : 32, gr = tq_ ? 2 : 1, gs = 4, ds_ = 4; // the down tile: Q2_0 8,1; else 4,2 - but on the M4 (no M5 layout) Q5_1 4,1 and Q8_0 8,1 (4,2 spills there: // 24-36 GB/s, mac-kern 513fc615) const bool m4 = !m5_layout(); - const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4, - dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1 || td.type == ref::Q5_0)) ? 1 : 2; + int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4, + dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1 || td.type == ref::Q5_0)) ? 1 : 2; const bool m16 = f16_enc_ && !stream_ && tg_.type == ref::Q2_0 && dq2 && !had_gu(l) && !had_dn(l); + // STRATA_METAL_TILES: the tiles measured here (calibrate_moe, after the window's first run); the fp16 experts keep theirs + const int dt = dq2 ? 142 : td.type; + bool moe_rec = false; + if (metal_tiles_file() && !stream_ && !m16 && moe_tunable(false, tg_.type) && moe_tunable(true, dt)) { + const auto gi = cal_moe_.find({0, tg_.type, ntok}), di = cal_moe_.find({1, dt, ntok}); + if (gi != cal_moe_.end() && di != cal_moe_.end()) { + const auto &g = gi->second, &d = di->second; + if (NFF % (g[2] * (32 / g[0]) * g[1]) == 0 && N % (d[2] * (32 / d[0]) * d[1]) == 0) { // another model's file: not its rows + gl = g[0]; gr = g[1]; gs = g[2]; + dl = d[0]; dr = d[1]; ds_ = d[2]; + } + } else moe_rec = cal_rec_ != nullptr; + } a.gx_gu = NFF / (gs * (32 / gl) * gr); a.gx_dn = N / (ds_ * (32 / dl) * dr); mtl::Dispatch r; @@ -1215,6 +1318,7 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 r.tg[0] = 512; } push(ds, r); + const int r_i = (int) ds.size() - 1; { GlueArgs ga{NFF, MAXT, N, TOPK}; mtl::Dispatch d; @@ -1254,6 +1358,7 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 g.ind_off = P_IGU * 4; if (cuts_) cuts_->push_back((int) ds.size()); // the host's turn: this layer's experts put in place push(ds, g); + const int g_i = (int) ds.size() - 1; nb_ = true; gemv(ds, T(b + "ffn_down_shexp.weight"), XH_, ysh_, N, ntok); if (hd) had_rows(ds, lib_, "d:", h_, 0, h_, 0, ntok * TOPK, NFF, hsh_.handle ? &hsh_ : nullptr, false); // H h, in place @@ -1268,6 +1373,8 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 dd.indirect = &plan_; dd.ind_off = P_IDN * 4; push(ds, dd); + if (moe_rec) + cal_rec_->moe.push_back({r_i, rf ? 8 : 4, g_i, (int) ds.size() - 1, tg_.type, dt, ntok, ct, {gl, gr, gs}, {dl, dr, ds_}}); } { GlueArgs a{N, (uint32_t) ntok, N, TOPK}; @@ -1401,6 +1508,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t seg_T_ = Tn; const double r0 = now_ms(); last_gpu_ms_ = 1000.0 * run_ds(c); + if (!c.moe.empty()) calibrate_moe(c); // its first run had no routing to read { static const int kt = std::getenv("STRATA_METAL_KPROF") ? std::atoi(std::getenv("STRATA_METAL_KPROF")) : 0; static int seen = 0; @@ -1415,6 +1523,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t rec_ = &cw.qsa; f16_enc_ = metal_f16() && head && !stream_ && !rows_; // decode windows of one resident sequence: not the prompt's, // not streaming, not batch slots + cal_rec_ = metal_calibrate_file() ? &cw : nullptr; cuts_ = stream_ ? &cw.cuts : nullptr; std::vector& ds = cw.ds; nb_ = false; @@ -1485,11 +1594,14 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t } rec_ = nullptr; f16_enc_ = false; + cal_rec_ = nullptr; cuts_ = nullptr; + if (!cw.gv.empty() || !cw.grd.empty()) calibrate(cw); seg_T_ = Tn; const double r0 = now_ms(); last_gpu_ms_ = 1000.0 * run_ds(cw); tp_rg_ = now_ms() - r0; + if (!cw.moe.empty()) calibrate_moe(cw); // the routing it reads is this run's last_n_disp_ = (int) ds.size(); return rows_ ? finish_rows(Tn, tokens, out, t0) : finish(Tn, tokens, pos0, head, out, t0); } @@ -1543,6 +1655,249 @@ bool Window::finish_rows(int Tn, const int32_t* tokens, int32_t* out, double t0) return true; } +// The tiles file: "device ", then "gemv type n_in n_out ntok L R S K", "grd kind ntok variant" and +// "moe gate/up=0|down=1 type ntok L R S" lines. +void Window::calibrate_load() { + if (metal_calibrate_file()) cal_file_ = metal_calibrate_file(); + FILE* f = std::fopen(metal_tiles_file(), "r"); + if (!f) { + if (!metal_calibrate_file()) + std::fprintf(stderr, "strata metal tiles: %s does not open: the defaults\n", metal_tiles_file()); + return; + } + char line[256] = {0}; + const std::string want = "device " + cx_->device_name(); + if (!std::fgets(line, sizeof line, f) || std::string(line).substr(0, std::strcspn(line, "\r\n")) != want) { + std::fprintf(stderr, "strata metal tiles: %s was measured on another GPU (not %s): the defaults\n", metal_tiles_file(), + cx_->device_name().c_str()); + std::fclose(f); + return; + } + char kind[16], pre[32]; + int a[8]; + while (std::fscanf(f, "%15s", kind) == 1) { + if (std::string(kind) == "gemv" && std::fscanf(f, "%d %d %d %d %d %d %d %d", a, a + 1, a + 2, a + 3, a + 4, a + 5, a + 6, a + 7) == 8) { + const std::array t{a[4], a[5], a[6], a[7]}; + if (gemv_tunable(a[0]) && std::find(std::begin(kGemvTiles), std::end(kGemvTiles), t) != std::end(kGemvTiles) && + gemv_tile_ok(t, a[0], a[1])) + cal_gemv_[{a[0], a[1], a[2], a[3]}] = t; + } else if (std::string(kind) == "grd" && std::fscanf(f, "%31s %d %d", pre, a, a + 1) == 3) { + if (a[1] >= 0 && a[1] < (int) std::size(kGrdVariants)) cal_grd_[{pre, a[0]}] = a[1]; + } else if (std::string(kind) == "moe" && std::fscanf(f, "%d %d %d %d %d %d", a, a + 1, a + 2, a + 3, a + 4, a + 5) == 6) { + const std::array t{a[3], a[4], a[5]}; + const bool down = a[0] == 1; + if (moe_tunable(down, a[1]) && (down ? std::find(std::begin(kDnTiles), std::end(kDnTiles), t) != std::end(kDnTiles) + : std::find(std::begin(kGuTiles), std::end(kGuTiles), t) != std::end(kGuTiles))) + cal_moe_[{a[0], a[1], a[2]}] = t; + } else break; + } + std::fclose(f); + std::fprintf(stderr, "strata metal tiles: %zu GEMV shapes, %zu gr_down kinds, %zu expert kernels from %s\n", + cal_gemv_.size(), cal_grd_.size(), cal_moe_.size(), metal_tiles_file()); +} + +// Each GEMV shape (and gr_down kind) of a new window, all its dispatches together (the layers' different weights, as +// the window reads them), timed on the candidate tiles: the best of 2 runs each, the outputs into a scratch buffer (the +// inputs are only read, so the model's state is not touched). The default stays unless a tile beats it by 3%. Every +// candidate for the decode windows (1 and 2 tokens; a larger window's GEMVs run at NTOK = 2 for that when the 2-token +// tile is not known yet - they read only the constant); a larger window then its default against the 2-token pick. +void Window::calibrate(Cached& c) { + const double t0 = now_ms(); + size_t need = 4 * (size_t) MAXT * (LR + HC); + for (const auto& g : c.gv) { + GemvArgs ga; + std::memcpy(&ga, c.ds[(size_t) g.i].args[4].bytes.data(), sizeof ga); + need = std::max(need, 4 * ((size_t) (g.shape[3] - 1) * ga.y_stride + (size_t) g.shape[2])); + } + mtl::Buffer scratch = cx_->alloc(need); + auto best_of = [&](const std::vector& ds) { + cx_->run(ds); + return std::min(cx_->run(ds), cx_->run(ds)); + }; + { // lift the GPU out of its idle clocks + std::vector w; + for (const auto& g : c.gv) { w.push_back(c.ds[(size_t) g.i]); w.back().args[3] = mtl::Arg::B(scratch); } + if (!w.empty()) cx_->warm(w, 0.3); + } + std::map, std::pair, std::vector>> gvs; + for (const auto& g : c.gv) { + auto& e = gvs[g.shape]; + e.first = g.tile; + e.second.push_back(g.i); + } + for (const auto& [shape, e] : gvs) { + const auto& [def, idx] = e; + const int ty = shape[0], n_in = shape[1], n_out = shape[2], ntok = shape[3]; + auto with = [&](const std::array& t, mtl::Dispatch d, int nt) { + const std::string fn = "gemv_" + std::to_string(ty) + "_" + std::to_string(t[0]) + "_" + std::to_string(t[1]) + "_" + + std::to_string(t[2]) + "_" + std::to_string(t[3]); + d.p = P(lib_, fn, mtl::Constants().i(0, nt), fn + "#" + std::to_string(nt)); + const int rpt = (t[2] / t[3]) * (32 / t[0]) * t[1]; + d.grid[0] = (uint32_t) ((n_out + rpt - 1) / rpt); + d.tg[0] = (uint32_t) (t[2] * 32); + return d; + }; + auto timed = [&](const std::array& t, int nt) { + std::vector ds; + for (int i : idx) { ds.push_back(with(t, c.ds[(size_t) i], nt)); ds.back().args[3] = mtl::Arg::B(scratch); } + return best_of(ds); + }; + // the better of `cands` against the default at nt (>= 3% faster, else the default) + auto pick = [&](const std::vector>& cands, int nt, double& td, double& tb) { + td = timed(def, nt); + std::array best = def; + tb = td; + for (const auto& t : cands) { + if (t == def || !gemv_tile_ok(t, ty, n_in)) continue; + const double tt = timed(t, nt); + if (tt < tb) { tb = tt; best = t; } + } + if (tb > 0.97 * td) { best = def; tb = td; } + return best; + }; + const int nd = std::min(ntok, 2); + const std::array kd{ty, n_in, n_out, nd}; + double td = 0, tb = 0; + if (!cal_gemv_.count(kd)) { + cal_gemv_[kd] = pick(std::vector>(std::begin(kGemvTiles), std::end(kGemvTiles)), nd, td, tb); + const auto& b = cal_gemv_[kd]; + std::fprintf(stderr, "strata metal calibrate: gemv type %d %d->%d T%d x%zu: %d,%d,%d,%d %.1f us -> %d,%d,%d,%d %.1f us\n", + ty, n_in, n_out, nd, idx.size(), def[0], def[1], def[2], def[3], 1e6 * td, b[0], b[1], b[2], b[3], 1e6 * tb); + } + if (ntok > 2) { + cal_gemv_[shape] = pick({cal_gemv_[kd]}, ntok, td, tb); + const auto& b = cal_gemv_[shape]; + std::fprintf(stderr, "strata metal calibrate: gemv type %d %d->%d T%d x%zu: %d,%d,%d,%d %.1f us -> %d,%d,%d,%d %.1f us\n", + ty, n_in, n_out, ntok, idx.size(), def[0], def[1], def[2], def[3], 1e6 * td, b[0], b[1], b[2], b[3], 1e6 * tb); + } + for (int i : idx) c.ds[(size_t) i] = with(cal_gemv_[shape], c.ds[(size_t) i], ntok); + } + std::map, std::vector> grs; + for (const auto& g : c.grd) grs[{g.pre, g.ntok}].push_back(g.i); + for (const auto& [k, idx] : grs) { + auto with = [&](int v, mtl::Dispatch d) { + const auto [sg, ks] = kGrdVariants[v]; + const std::string fn = k.first + std::to_string(sg) + "_" + std::to_string(ks); + d.p = P(lib_, fn, mtl::Constants().i(1, k.second), fn + "#" + std::to_string(k.second)); + d.grid[0] = (uint32_t) ((LR + HC + sg / ks - 1) / (sg / ks)); + d.tg[0] = (uint32_t) (32 * sg); + return d; + }; + auto timed = [&](int v) { + std::vector ds; + for (int i : idx) { + ds.push_back(with(v, c.ds[(size_t) i])); + ds.back().args[3] = mtl::Arg::B(scratch); // lo + ds.back().args[4] = mtl::Arg::B(scratch, 4 * (size_t) MAXT * LR); // the inject values + } + return best_of(ds); + }; + const double td = timed(0); + int best = 0; + double tb = td; + for (int v = 1; v < (int) std::size(kGrdVariants); ++v) { + const mtl::Dispatch d = with(v, c.ds[(size_t) idx[0]]); + if (d.p->max_threads < (int) d.tg[0]) continue; // above this pipeline's threadgroup limit (require_fit) + const double tt = timed(v); + if (tt < tb) { tb = tt; best = v; } + } + if (tb > 0.97 * td) { best = 0; tb = td; } + cal_grd_[k] = best; + for (int i : idx) c.ds[(size_t) i] = with(best, c.ds[(size_t) i]); + std::fprintf(stderr, "strata metal calibrate: %s T%d x%zu: 8_8 %.1f us -> %d_%d %.1f us\n", k.first.c_str(), k.second, + idx.size(), 1e6 * td, kGrdVariants[best][0], kGrdVariants[best][1], 1e6 * tb); + } + cx_->free(scratch); + c.gv.clear(); + c.grd.clear(); + calibrate_save(); + std::fprintf(stderr, "strata metal calibrate: measured in %.0f ms, kept in %s\n", now_ms() - t0, cal_file_.c_str()); +} + +void Window::calibrate_save() { + FILE* f = cal_file_.empty() ? nullptr : std::fopen(cal_file_.c_str(), "w"); + if (!f) return; + std::fprintf(f, "device %s\n", cx_->device_name().c_str()); + for (const auto& [s, t] : cal_gemv_) std::fprintf(f, "gemv %d %d %d %d %d %d %d %d\n", s[0], s[1], s[2], s[3], t[0], t[1], t[2], t[3]); + for (const auto& [k, v] : cal_grd_) std::fprintf(f, "grd %s %d %d\n", k.first.c_str(), k.second, v); + for (const auto& [k, t] : cal_moe_) std::fprintf(f, "moe %d %d %d %d %d %d\n", k[0], k[1], k[2], t[0], t[1], t[2]); + std::fclose(f); +} + +// The routed experts of a window that has just run: the layers of each pair of formats (a pack mixes them) together, +// their gate/up (then down) dispatches on the routing the last layer left in plan_ (real experts of each layer, a real +// count of them), with a direct grid in place of the indirect one - moe_route writes gx into the plan, the kernels do +// not read it. The outputs go to a scratch buffer. +void Window::calibrate_moe(Cached& c) { + const int og = plan_.as()[P_IGU + 1]; + if (og <= 0) return; // no routing yet: the next run + const double t0 = now_ms(); + mtl::Buffer scratch = cx_->alloc(std::max(h_.bytes, ye_.bytes)); + std::map, std::vector> fmts; // (gate/up format, down format) -> its layers + for (const auto& m : c.moe) fmts[{m.gt, m.dt}].push_back(&m); + for (const auto& [fmt, layers] : fmts) { + const Cached::Moe& m0 = *layers.front(); + for (int down = 0; down < 2; ++down) { + const int ty = down ? fmt.second : fmt.first, rows = down ? N : NFF; + const std::array def = down ? m0.dn : m0.gu; + auto with = [&](const std::array& t, mtl::Dispatch d) { + const std::string fn = std::string(down ? "moe_down_" : "moe_gu_") + std::to_string(ty) + "_" + std::to_string(t[0]) + + "_" + std::to_string(t[1]) + "_" + std::to_string(t[2]); + d.p = P(lib_, fn, mtl::Constants().b(23, m0.ct), fn + (m0.ct ? "#ct" : "")); + d.tg[0] = (uint32_t) (32 * t[2]); + return d; + }; + auto timed = [&](const std::array& t) { + std::vector ds; + for (const Cached::Moe* m : layers) { + ds.push_back(with(t, c.ds[(size_t) (down ? m->d : m->g)])); + mtl::Dispatch& d = ds.back(); + d.indirect = nullptr; + d.grid[0] = (uint32_t) (rows / (t[2] * (32 / t[0]) * t[1])); + d.grid[1] = (uint32_t) og; + d.grid[2] = 1; + d.args[2] = mtl::Arg::B(scratch); // gate/up's h, down's ye + } + cx_->run(ds); + return std::min(cx_->run(ds), cx_->run(ds)); + }; + std::array best = def; + double td = 0, tb = 0; + const auto known = cal_moe_.find({down, ty, m0.ntok}); // measured meanwhile by another window + if (known != cal_moe_.end()) best = known->second; + else { + td = tb = timed(def); + const auto* b0 = down ? std::begin(kDnTiles) : std::begin(kGuTiles); + const auto* b1 = down ? std::end(kDnTiles) : std::end(kGuTiles); + for (const auto* t = b0; t != b1; ++t) { + if (*t == def || rows % ((*t)[2] * (32 / (*t)[0]) * (*t)[1]) != 0) continue; + const double tt = timed(*t); + if (tt < tb) { tb = tt; best = *t; } + } + if (tb > 0.97 * td) { best = def; tb = td; } + cal_moe_[{down, ty, m0.ntok}] = best; + } + const int gx = rows / (best[2] * (32 / best[0]) * best[1]); + for (const Cached::Moe* m : layers) { + mtl::Dispatch& d = c.ds[(size_t) (down ? m->d : m->g)]; + d = with(best, d); + MoeArgs a; // the route writes gx into the plan: its copy of the arguments + std::vector& rb = c.ds[(size_t) m->r].args[(size_t) m->ra].bytes; + std::memcpy(&a, rb.data(), sizeof a); + (down ? a.gx_dn : a.gx_gu) = (uint32_t) gx; + std::memcpy(rb.data(), &a, sizeof a); + } + std::fprintf(stderr, "strata metal calibrate: moe_%s type %d T%d x%zu (%d experts): %d,%d,%d %.1f us -> %d,%d,%d %.1f us\n", + down ? "down" : "gu", ty, m0.ntok, layers.size(), og, def[0], def[1], def[2], 1e6 * td, best[0], best[1], + best[2], 1e6 * tb); + } + } + cx_->free(scratch); + c.moe.clear(); + calibrate_save(); + std::fprintf(stderr, "strata metal calibrate: experts measured in %.0f ms\n", now_ms() - t0); +} + double Window::run_ds(Cached& c) { if (!stream_) return cx_->run_graph(c.ds); return cx_->run_graph_seg(c.ds, c.cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this); diff --git a/strata_control/calibration.py b/strata_control/calibration.py index 075551fea36..c9907e404d3 100644 --- a/strata_control/calibration.py +++ b/strata_control/calibration.py @@ -275,6 +275,8 @@ def measure(ref: str, output: Path, *, exclusive_ack: bool = False, say=print, r if not exclusive_ack: raise ProfileError("calibration measure requires an exclusive-access acknowledgement") path, profile, revision, cfg = _target(ref) + if cfg.get("backend") == "metal": # its tiles replace the file the config names: setup.sh's review only + raise ProfileError("a Mac calibrates with ./setup.sh --calibrate", 2) with profiles._lock(_lock_name(path)), _managed_exclusive(): current_path, profile, revision, cfg = _target(ref) if current_path != path or (expected_profile_revision is not None and diff --git a/tests/terminal/test_calibration.py b/tests/terminal/test_calibration.py index 90c927696cb..deb06e4edcb 100644 --- a/tests/terminal/test_calibration.py +++ b/tests/terminal/test_calibration.py @@ -333,6 +333,14 @@ def test_missing_report_folder_is_refused_before_the_sweep(self): calibration.measure("local", output, exclusive_ack=True, runner=lambda *_a, **_k: self.fail("the sweep must not start")) + def test_a_mac_config_is_refused(self): + # a Mac's tiles go through ./setup.sh --calibrate: measure would replace the file the config names + cfg = json.loads(self.config.read_text()) + self.config.write_text(json.dumps({**cfg, "backend": "metal"})) + with self.assertRaisesRegex(profiles.ProfileError, "setup.sh --calibrate"): + calibration.measure("local", self.root / "report.json", exclusive_ack=True, + runner=lambda *_a, **_k: self.fail("the sweep must not start")) + def test_calibration_rechecks_reviewed_digests_inside_operation_lock(self): report_path = self.root / "pinned-calibration.json" with self.assertRaisesRegex(profiles.ProfileError, "source changed after review"): diff --git a/tools/calibrate.py b/tools/calibrate.py index f2408ae0d79..4d13fdf2999 100644 --- a/tools/calibrate.py +++ b/tools/calibrate.py @@ -18,6 +18,13 @@ magnitude between rigs (about 700 tokens on four GP100 dies, 5000 on two RTX 3090 Ti) and costs when it is wrong (a fixed 1024 read 22-36% slower on two RTX 4090) - so it is measured here, once, as the prompt read speed of fresh prompts, rather than guessed or learned while the PC serves. + STRATA_METAL_TILES (Mac only) + the tiles of the GPU kernels a decode window runs (the dense GEMVs, the hyper-connection read, the + routed experts). The engine's defaults were tuned on the M4 and M5; a GPU with less ALU per byte of + bandwidth wants others (the M1 Max: a 2-token window's GEMVs 33 -> 21 ms). The engine measures them + itself (STRATA_METAL_CALIBRATE) while it answers a few prompts, then they race the defaults and go + into the config's "env" only when they win. The answers may differ slightly from the defaults' (the + order of a sum: a near-tie can flip). STRATA_SH_STREAM (AMD only) the shared expert on its own stream. Off by default on HIP (#816: an RX 6800 on Windows and 4x R9700 decode 12-45% faster without it), but the same card family on another PC goes the other way (#1272: @@ -36,6 +43,7 @@ from __future__ import annotations import json +import re import statistics import sys import threading @@ -55,6 +63,8 @@ # HIP: the shared-expert stream fork, off (the HIP default) against on; an engine environment variable, not an argument SH_STREAM = "STRATA_SH_STREAM" ADAPT_CANDIDATES = (None, ("1", "80", "0.97"), ("1", "160", "0.97")) +# a Mac: the engine measures its kernels' tiles into CALIBRATE's file; TILES points the engine at the result +CALIBRATE, TILES = "STRATA_METAL_CALIBRATE", "STRATA_METAL_TILES" ADAPT_FLAGS = ("--adapt-every", "--adapt-swaps", "--adapt-decay") MAX_NEW = 128 # the expert tier follows a text over some windows: 128-token answers end before it shows (on a Xeon E5-2673 v3 with @@ -240,6 +250,10 @@ def start_engine(args, env=None): tok = ST.Tokenizer(toks, (tpath / "merges.txt").read_text(encoding="utf-8").split("\n"), json.loads((tpath / "token_type.json").read_text())) ids_list = [chat_ids(tok, p) for p in PROMPTS] + if sys.platform == "darwin": # a file a model: another model's calibration must not take this one's shapes away + name = re.sub(r"[^A-Za-z0-9._-]", "-", cfg.get("model_name") or "model") + return measure_metal(engine_args(cfg), ids_list, start_engine, say, Path(cfg.get("cwd") or ROOT) / f"metal-tiles-{name}.txt", + progress=progress, cancel_event=cancel_event) if progress is None and cancel_event is None: return measure(engine_args(cfg), ids_list, start_engine, say, host_worker_extras(), doc_prompt(tok), hip=cfg.get("backend") == "hip") @@ -446,6 +460,50 @@ def fork_is_unsafe(args) -> bool: return sys.platform == "win32" and arg_value(args, "--kv-resident") is not None +def measure_metal(base_args: list[str], ids_list, start_engine, say=print, tiles: Path = ROOT / "metal-tiles.txt", *, + progress=None, cancel_event=None) -> dict: + """A Mac: the engine measures its decode kernels' tiles on this GPU while it answers the prompts, then the tiles race + the defaults - a restart each, interleaved rounds, on the config's own arguments (the PC's product defaults are not + the Mac's) - and are kept (env STRATA_METAL_TILES) only when they win by more than MIN_GAIN. Both variables are set + on every start (empty: off), so a tiles file the config names already does not decide the baseline.""" + t0 = time.time() + report: dict = {} + new = tiles.with_suffix(".new.txt") + new.unlink(missing_ok=True) + _check_cancel(cancel_event) + _progress(progress, "loading", 0, 1) + say(" Measuring the GPU kernels' tiles on this Mac ...") + e = start_engine(base_args, {CALIBRATE: str(new), TILES: ""}) + try: + Session(e, ids_list, cancel_event).warm_up(1) + finally: + close(e) + _progress(progress, "loading", 1, 1) + if not new.exists(): + raise RuntimeError("the engine wrote no tiles") + rates: dict = {"default": [], "tiles": []} + for rnd in range(SWEEP_ROUNDS): + for key in ("default", "tiles") if rnd % 2 == 0 else ("tiles", "default"): + _check_cancel(cancel_event) + env = {CALIBRATE: "", TILES: str(new) if key == "tiles" else ""} + say(f" Measuring the {'measured tiles' if key == 'tiles' else 'default tiles'} (restarts the engine) ...") + got: dict = {} + measure_candidate(lambda a: start_engine(a, env), base_args, key, say, report, got, key, ids_list, + cancel_event=cancel_event) + rates[key] += got.get(key, []) + report["metal_tiles"] = rates + env = {} + if pick(rates, "default") == "tiles": + new.replace(tiles) + env[TILES] = str(tiles) + else: + new.unlink(missing_ok=True) + best = rates["tiles" if env else "default"] + report["seconds"] = round(time.time() - t0) + report["tok_s"] = round(statistics.median(best), 1) if best else None + return {"settings": {}, "env": env, "report": report} + + def restart(start_engine, args, what, say, report): """#1337: a later step's engine start that fails is a candidate that loses, not a failed run: what the earlier steps measured (and the settings they chose) must survive it. Returns the engine, or None after saying why and @@ -511,7 +569,7 @@ def apply(args: list[str], settings: dict) -> list[str]: return out -ENV_KEYS = (SH_STREAM,) +ENV_KEYS = (SH_STREAM, TILES, CALIBRATE) def apply_env(env: dict | None, calibrated: dict | None) -> dict: diff --git a/tools/test_calibrate.py b/tools/test_calibrate.py index 7ffe1fb63bb..e42a67eb43b 100644 --- a/tools/test_calibrate.py +++ b/tools/test_calibrate.py @@ -292,6 +292,58 @@ def start(args, env=None): self.assertNotIn("sh_stream", cuda["report"]) self.assertTrue(all(e == {} for e in envs)) + def test_metal_tiles(self): + # a Mac: the engine measures its tiles (STRATA_METAL_CALIBRATE) on the first start, then the tiles race the + # defaults with a restart each, both variables set explicitly; kept only when they win by more than MIN_GAIN + mac = ["--pack", "p", "--spec", "2", "--max-context", "8192"] # a Mac config: no PC product defaults + + def run(bonus, d): + envs = [] + + def start(args, env=None): + self.assertEqual(args, mac) # the config's own arguments, every start + envs.append(dict(env or {})) + if (env or {}).get(CAL.CALIBRATE): + Path(env[CAL.CALIBRATE]).write_text("device Apple M1 Max\n") + tiles = bool((env or {}).get(CAL.TILES)) + return FakeEngine(args, lambda f, p, w: 50.0 * (bonus if tiles else 1.0), 6) + res = CAL.measure_metal(mac, [[1, 2, 3]] * 3, start, say=lambda *_: None, tiles=Path(d) / "metal-tiles.txt") + return res, envs + with tempfile.TemporaryDirectory() as d: + res, envs = run(1.20, d) + tiles = str(Path(d) / "metal-tiles.txt") + self.assertEqual(res["env"], {CAL.TILES: tiles}) + self.assertTrue(Path(tiles).exists() and not Path(d, "metal-tiles.new.txt").exists()) + self.assertEqual(envs[0][CAL.TILES], "") # the measuring start reads no older tiles + self.assertTrue(all(set(e) == {CAL.CALIBRATE, CAL.TILES} for e in envs)) # never left to the config's env + self.assertEqual(sum(1 for e in envs[1:] if e[CAL.TILES]), CAL.SWEEP_ROUNDS) + self.assertGreater(res["report"]["tok_s"], 55.0) # the reported speed is the tiles' + Path(tiles).unlink() + flat, _ = run(1.02, d) # within noise: the defaults stay, no file + self.assertEqual(flat["env"], {}) + self.assertFalse(Path(tiles).exists() or Path(d, "metal-tiles.new.txt").exists()) + + def test_engine_death_during_a_metal_round_keeps_the_run(self): + # the tiles' second engine exits mid-measurement: that round is dropped, the other rounds decide + mac = ["--pack", "p", "--spec", "2", "--max-context", "8192"] + tile_starts = [] + + def start(args, env=None): + if (env or {}).get(CAL.CALIBRATE): + Path(env[CAL.CALIBRATE]).write_text("device Apple M1 Max\n") + tiles = bool((env or {}).get(CAL.TILES)) + e = FakeEngine(args, lambda f, p, w: 50.0 * (1.20 if tiles else 1.0), 6) + if tiles: + tile_starts.append(1) + e.dies = len(tile_starts) == 2 + return e + with tempfile.TemporaryDirectory() as d: + tiles = str(Path(d) / "metal-tiles.txt") + res = CAL.measure_metal(mac, [[1, 2, 3]] * 3, start, say=lambda *_: None, tiles=Path(tiles)) + self.assertEqual(res["env"], {CAL.TILES: tiles}) + self.assertEqual(len(res["report"]["metal_tiles"]["tiles"]), 2 * (CAL.SWEEP_ROUNDS - 1)) + self.assertIn("tiles", res["report"]["failed_starts"]) + def test_a_failed_fork_start_keeps_the_default(self): def start(args, env=None): if (env or {}).get(CAL.SH_STREAM) == "1": @@ -318,6 +370,8 @@ def test_apply_env(self): self.assertEqual(CAL.apply_env(old, {}), {"STRATA_ROCBLAS_TUNING": "tools/hip"}) # an older choice goes self.assertEqual(CAL.apply_env({"X": "1"}, {CAL.SH_STREAM: "1"}), {"X": "1", CAL.SH_STREAM: "1"}) self.assertEqual(CAL.apply_env(None, None), {}) + mac = {"STRATA_METAL_STREAM": "0", "STRATA_METAL_F16": "1"} # a Mac's own variables stay, an older tiles file goes + self.assertEqual(CAL.apply_env({**mac, CAL.TILES: "old.txt"}, {}), mac) def test_apply(self): a = CAL.apply(BASE, {"--pcie-frac": "0.35", "--pool-workers": "4"}) @@ -354,7 +408,7 @@ def test_run_measures_with_the_servers_args(self): seen.append((args, hip)) or {}) fake = type("ST", (), {"Tokenizer": lambda *a: type("T", (), {"encode": lambda s, t, **k: [0]})()}) try: - with mock.patch.dict(sys.modules, {"strata_tokenizer": fake}): + with mock.patch.dict(sys.modules, {"strata_tokenizer": fake}), mock.patch.object(CAL.sys, "platform", "linux"): CAL.run({"tokenizer": d, "args": list(BASE), "gpu": [0]}, start_engine=lambda a: None) CAL.run({"tokenizer": d, "args": list(BASE), "gpu": [0], "backend": "hip"}, start_engine=lambda a: None) finally: