| """Measure what changing the number of active MoE experts actually does. |
| |
| Run against any MoE GGUF. It loads the model once per expert count, times a |
| fixed generation, and reports throughput. The point is to replace a guess with |
| a number before touching a serving configuration. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import time |
| from pathlib import Path |
|
|
| from llama_cpp import Llama |
|
|
| PROMPT = "Write a Python function that reverses a string." |
|
|
|
|
| def measure( |
| model_path: str, |
| expert_count: int | None, |
| tokens: int, |
| threads: int, |
| ) -> dict: |
| """Load the model with an optional expert override and time a generation.""" |
| key = "expert_used_count" |
| overrides = None |
| if expert_count is not None: |
| overrides = {key: expert_count} |
|
|
| kwargs = {"kv_overrides": overrides} if overrides else {} |
| started = time.monotonic() |
| try: |
| model = Llama( |
| model_path=model_path, |
| n_ctx=512, |
| n_threads=threads, |
| verbose=False, |
| **kwargs, |
| ) |
| except Exception as exc: |
| return {"experts": expert_count, "error": f"{type(exc).__name__}: {exc}"} |
| load_s = time.monotonic() - started |
|
|
| arch = model.metadata.get("general.architecture", "") |
| native = model.metadata.get(f"{arch}.{key}") |
|
|
| started = time.monotonic() |
| out = model.create_chat_completion( |
| messages=[{"role": "user", "content": PROMPT}], |
| max_tokens=tokens, |
| temperature=0.0, |
| ) |
| gen_s = time.monotonic() - started |
| produced = out.get("usage", {}).get("completion_tokens", 0) |
| text = out["choices"][0]["message"]["content"] |
| del model |
|
|
| return { |
| "experts": expert_count, |
| "native_experts": native, |
| "load_s": round(load_s, 2), |
| "gen_s": round(gen_s, 2), |
| "tokens": produced, |
| "tok_per_s": round(produced / gen_s, 1) if gen_s else None, |
| "preview": text[:120], |
| } |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("model", help="path to a MoE GGUF file") |
| parser.add_argument( |
| "--experts", |
| default="2,8,15", |
| help="comma-separated expert counts to try (omit for native only)", |
| ) |
| parser.add_argument("--tokens", type=int, default=128) |
| parser.add_argument("--threads", type=int, default=4) |
| parser.add_argument("--json", help="write the results to this JSON file") |
| args = parser.parse_args() |
|
|
| if not Path(args.model).is_file(): |
| parser.error(f"model not found: {args.model}") |
|
|
| counts: list[int | None] = [None] |
| counts += [int(c) for c in args.experts.split(",") if c.strip()] |
|
|
| results = [] |
| for count in counts: |
| label = "native" if count is None else f"{count} experts" |
| print(f"--- {label}", flush=True) |
| row = measure(args.model, count, args.tokens, args.threads) |
| results.append(row) |
| if "error" in row: |
| print(f" FAILED: {row['error']}", flush=True) |
| else: |
| print( |
| f" {row['tok_per_s']} tok/s " |
| f"(load {row['load_s']}s, gen {row['gen_s']}s)", |
| flush=True, |
| ) |
|
|
| baseline = next((r for r in results if r.get("experts") is None), None) |
| if baseline and baseline.get("tok_per_s"): |
| print("\nrelative to native:", flush=True) |
| for row in results: |
| if row.get("tok_per_s"): |
| ratio = row["tok_per_s"] / baseline["tok_per_s"] |
| print( |
| f" {str(row['experts']):>4} experts -> x{ratio:.2f} speed", |
| flush=True, |
| ) |
|
|
| if args.json: |
| Path(args.json).write_text( |
| json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8" |
| ) |
| print(f"\nwritten to {args.json}", flush=True) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
