"""Command-line entry points. Imports optional HTTP dependencies only for serve.""" from __future__ import annotations import argparse import json import statistics import sys import time def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(prog="cism") commands = parser.add_subparsers(dest="command", required=True) for name in ("inspect", "generate", "bench", "cache-report"): command = commands.add_parser(name) command.add_argument("model", help="HF repository ID or local safetensors directory") command.add_argument("--precision", choices=["fp32", "int8", "hybrid-int4", "hybrid-fp4"], default="fp32") command.add_argument("--act-precision", choices=["fp32", "int8"], default="fp32", help="activation precision for quantized matrices (default fp32)") command.add_argument("--revision") command.add_argument("--trust-remote-code", action="store_true") command.add_argument("--local-files-only", action="store_true") command.add_argument("--lock-pages", action="store_true", help="pin weight pages in physical RAM (VirtualLock/mlock)") command.add_argument("--threads", type=int, default=1, metavar="N", help="native decode worker threads, including caller (1-64)") if name in ("generate", "bench"): command.add_argument("--prompt", default="The future of artificial intelligence is") command.add_argument("--max-tokens", type=int, default=64) if name == "generate": command.add_argument("--temperature", type=float, default=0) command.add_argument("--seed", type=int, default=0) command.add_argument("--spec", type=int, default=0, metavar="K", help="greedy prompt-lookup speculation (2-16, temperature must be 0)") if name == "bench": command.add_argument("--runs", type=int, default=5, help="measured repetitions (legacy name; --reps overrides when given)") command.add_argument("--reps", type=int, default=None, help="measured repetitions per setup; median reported (default: --runs value)") command.add_argument("--spec", type=int, default=0, metavar="K", help="greedy prompt-lookup speculation (2-16)") if name == "cache-report": command.add_argument("--decode-tokens", type=int, default=64) command.add_argument("--scrub-mb", type=int, default=96, help="cache-eviction buffer size in MiB (use ~2x L3)") serve = commands.add_parser("serve") serve.add_argument("models", nargs="+", help="one or more HF repository IDs or local safetensors directories") serve.add_argument("--precision", choices=["fp32", "int8", "hybrid-int4", "hybrid-fp4"], default="fp32") serve.add_argument("--act-precision", choices=["fp32", "int8"], default="fp32", help="activation precision for quantized matrices (default fp32)") serve.add_argument("--revision") serve.add_argument("--trust-remote-code", action="store_true") serve.add_argument("--local-files-only", action="store_true") serve.add_argument("--lock-pages", action="store_true", help="pin weight pages in physical RAM (VirtualLock/mlock)") serve.add_argument("--threads", type=int, default=1, metavar="N", help="native decode worker threads per model, including caller (1-64)") serve.add_argument("--keep-warm", type=float, default=0.0, metavar="SECONDS", help="touch weights of every loaded model every SECONDS while idle") serve.add_argument("--allow-remote-code-load", action="store_true", help="permit POST /v1/models/load with trust_remote_code=true") serve.add_argument("--host", default="127.0.0.1") serve.add_argument("--port", type=int, default=8000) serve.add_argument("--queue-size", type=int, default=8) tui = commands.add_parser("tui") tui.add_argument("model", help="HF repository ID or local safetensors directory") tui.add_argument("--precision", choices=["fp32", "int8", "hybrid-int4", "hybrid-fp4"], default="int8", help="weight precision (default int8)") tui.add_argument("--threads", type=int, default=2, metavar="N", help="native decode worker threads, including caller (1-64)") tui.add_argument("--max-tokens", type=int, default=128, metavar="N", help="max new tokens per generation (>=1)") tui.add_argument("--temperature", type=float, default=0.0, help="sampling temperature, finite >=0 (0 for greedy)") tui.add_argument("--top-p", type=float, default=1.0, help="nucleus sampling mass in (0, 1]") tui.add_argument("--top-k", type=int, default=0, metavar="K", help="top-k cutoff >=0 (0 disables)") tui.add_argument("--seed", type=int, default=0, metavar="S", help="sampling seed in [0, 2**64)") tui.add_argument("--spec-k", type=int, default=0, metavar="K", help="greedy prompt-lookup speculation (0 off or 2-16, requires temperature 0)") tui.add_argument("--trust-remote-code", action="store_true", help="allow remote code when loading the model") tui.add_argument("--local-files-only", action="store_true", help="load only from local cache, never download") args = parser.parse_args(argv) if args.command == "bench": _reps = args.reps if args.reps is not None else args.runs if _reps < 1 or args.runs < 1 or args.max_tokens < 2: parser.error("bench requires --runs/--reps >= 1 and --max-tokens >= 2") if args.command == "tui": import math as _math if not 1 <= args.threads <= 64: parser.error("tui requires --threads in [1, 64]") if args.max_tokens < 1: parser.error("tui requires --max-tokens >= 1") if not _math.isfinite(args.temperature) or args.temperature < 0: parser.error("tui requires --temperature finite and >= 0") if not _math.isfinite(args.top_p) or not 0 < args.top_p <= 1: parser.error("tui requires --top-p in (0, 1]") if args.top_k < 0: parser.error("tui requires --top-k >= 0") if not 0 <= args.seed < 2**64: parser.error("tui requires --seed in [0, 2**64)") if args.spec_k != 0 and not 2 <= args.spec_k <= 16: parser.error("tui requires --spec-k 0 or in [2, 16]") if args.spec_k >= 2 and args.temperature != 0: parser.error("tui requires --temperature 0 when --spec-k is used") from .tui import run as _tui_run class _TeeStdout: """Mirror tui 'error: ...' lines to stderr for CLI conventions.""" def __init__(self, out, err): self._out = out self._err = err def write(self, s): try: res = self._out.write(s) except Exception: res = None try: if isinstance(s, str) and s.startswith("error:"): self._err.write(s) try: self._err.flush() except Exception: pass except Exception: pass if res is not None: return res return len(s) if isinstance(s, str) else 0 def flush(self): try: self._out.flush() except Exception: pass def __getattr__(self, name): return getattr(self._out, name) _real_out, _real_err = sys.stdout, sys.stderr sys.stdout = _TeeStdout(_real_out, _real_err) try: try: return _tui_run( args.model, precision=args.precision, threads=args.threads, max_new_tokens=args.max_tokens, temperature=args.temperature, top_p=args.top_p, top_k=args.top_k, seed=args.seed, spec_k=args.spec_k, trust_remote_code=args.trust_remote_code, local_files_only=args.local_files_only, ) except (ValueError, RuntimeError, ImportError, OSError) as exc: print(f"error: {exc}", file=_real_err) return 1 finally: sys.stdout = _real_out from .engine import Engine try: engines = {} for model in (args.models if args.command == "serve" else [args.model]): engine = Engine.from_pretrained(model, precision=args.precision, revision=args.revision, trust_remote_code=args.trust_remote_code, local_files_only=args.local_files_only, threads=args.threads, act_precision=args.act_precision) if args.lock_pages and engine.lock_pages() == 0: print("cism: page locking failed for this process; continuing unlocked", file=sys.stderr) engines[engine.model_id] = engine if args.command == "inspect": print(json.dumps(next(iter(engines.values())).info, indent=2)) elif args.command == "generate": engine = next(iter(engines.values())) for chunk in engine.stream(args.prompt, max_new_tokens=args.max_tokens, temperature=args.temperature, seed=args.seed, spec_k=args.spec): print(chunk.text, end="", flush=True) print() elif args.command == "cache-report": import numpy as np engine = next(iter(engines.values())) weight_bytes = engine.info["weight_bytes"] def timed_scan(repeats: int, evictions: int) -> float: # Unique 2x-L3 writes evict residency; the scan is timed alone. scrub = np.empty(args.scrub_mb * 1024 * 1024, dtype=np.uint8) samples = [] for _ in range(repeats): for _ in range(evictions): scrub[:] = 1 start = time.perf_counter() engine.scan() samples.append(time.perf_counter() - start) return statistics.median(samples) warm_seconds = timed_scan(20, evictions=0) cold_seconds = timed_scan(5, evictions=2) prompt = engine.encode("cache report probe") session = engine._native.create_session(prompt, max_new_tokens=args.decode_tokens, eos_token_ids=[]) session.next_tokens(1) start = time.perf_counter() generated = session.next_tokens(args.decode_tokens - 1) decode_seconds = time.perf_counter() - start tokens_per_second = len(generated) / decode_seconds if generated and decode_seconds > 0 else 0 effective_gbps = weight_bytes * tokens_per_second / 1e9 warm_gbps = weight_bytes / warm_seconds / 1e9 cold_gbps = weight_bytes / cold_seconds / 1e9 ratio = cold_seconds / warm_seconds if warm_seconds > 0 else 0 if ratio >= 2: residency = "cache-resident when warm (cold scan is much slower)" elif ratio >= 1.3: residency = "partially cache-resident" else: residency = "no measurable residency benefit from this probe" if not tokens_per_second: bottleneck = "decode produced no tokens" elif effective_gbps <= cold_gbps: bottleneck = ("compute-bound: decode bandwidth is below even the cold DRAM scan; " "the cache is not the bottleneck, the kernels are") elif effective_gbps < warm_gbps * 0.8: bottleneck = "DRAM-bandwidth-limited: weights are effectively streaming from RAM" else: bottleneck = "bandwidth-limited at cache-resident rates" print(json.dumps({**engine.info, "weight_bytes": weight_bytes, "warm_scan_gbps": round(warm_gbps, 2), "cold_scan_gbps": round(cold_gbps, 2), "warm_to_cold_ratio": round(ratio, 2), "decode_tokens_per_second": round(tokens_per_second, 1), "decode_effective_gbps": round(effective_gbps, 2), "residency": residency, "decode_bottleneck": bottleneck}, indent=2)) elif args.command == "bench": from .autobench import _cpu_snapshot, _thermal_status engine = next(iter(engines.values())) prompt = engine.encode(args.prompt) reps = args.reps if args.reps is not None else args.runs # Deliberately ignore EOS for fixed-length comparison; report this. ttfts, decode_rates = [], [] for run in range(reps + 1): start = time.perf_counter() session = engine._native.create_session(prompt, max_new_tokens=args.max_tokens, eos_token_ids=[]) first = session.next_tokens(1) after_first = time.perf_counter() tail = session.next_tokens(args.max_tokens - 1) end = time.perf_counter() if len(first) + len(tail) != args.max_tokens: raise RuntimeError("Benchmark failed to generate the requested token count") if run: ttfts.append((after_first - start) * 1000) decode_rates.append(len(tail) / (end - after_first)) throttled, drop_percent = _thermal_status(decode_rates) try: cpu_snapshot = _cpu_snapshot() except Exception: cpu_snapshot = {} if not isinstance(cpu_snapshot, dict): cpu_snapshot = {} print(json.dumps({**engine.info, "prompt_tokens": len(prompt), "generated_tokens": args.max_tokens, "runs": reps, "reps": reps, "eos_ignored": True, "warmup_runs": 1, "median_prefill_first_token_ms": statistics.median(ttfts), "median_decode_tokens_per_second": statistics.median(decode_rates), "decode_tokens_per_second": decode_rates, "decode_tokens_per_second_min": min(decode_rates), "decode_tokens_per_second_max": max(decode_rates), "decode_tokens_per_second_std": round(statistics.pstdev(decode_rates), 3) if len(decode_rates) >= 2 else 0.0, "prefill_first_token_ms_min": min(ttfts), "prefill_first_token_ms_max": max(ttfts), "thermal_throttled": throttled, "thermal_drop_percent": drop_percent, "cpu_snapshot": cpu_snapshot}, indent=2)) else: import uvicorn from .server import create_app if args.keep_warm: for engine in engines.values(): engine.start_keep_warm(args.keep_warm) loaded = engines if len(engines) > 1 else next(iter(engines.values())) uvicorn.run(create_app(loaded, queue_size=args.queue_size, allow_remote_code_load=args.allow_remote_code_load, load_lock_pages=args.lock_pages, load_keep_warm=args.keep_warm, load_threads_default=args.threads, load_local_files_only=args.local_files_only, load_act_precision=args.act_precision), host=args.host, port=args.port, workers=1) except (ValueError, RuntimeError, ImportError, OSError) as exc: print(f"cism: {exc}", file=sys.stderr) return 1 return 0