diff --git a/inference_engine/bridge/manifest.py b/inference_engine/bridge/manifest.py index a93ef592..a9d31542 100644 --- a/inference_engine/bridge/manifest.py +++ b/inference_engine/bridge/manifest.py @@ -650,6 +650,36 @@ def _harness_preset( }, validate_reports=False, ), + Preset( + name="vllm-mlx-niah", + description="Evaluate vLLM-MLX (Apple-Silicon continuous batching) on " + "the SAME local MLX gemma verifier used for the B>1,L=1 " + "batched-decode bug repro: pip install vllm-mlx, serve with " + "--continuous-batching, fire N concurrent NIAH requests " + "(each with a unique needle), report per-session recall + " + "aggregate decode tok/s vs N=1. Answers whether vLLM-MLX is " + "BOTH parallel AND recall-preserving on our config — the " + "Mac analog of the CUDA KIE-v2 (Kakeya Attention on a " + "borrowed runtime).", + command_templates=( + ("python3", "-m", "pip", "install", "--upgrade", "vllm-mlx"), + ( + "python3", "scripts/research/vllm_mlx_niah_bench.py", + "--model-path", "${ENV:KAKEYA_MAC_VERIFIER_PATH}", + "--sessions", "{n_samples}", + "--haystack-lines", "60", + "--max-new-tokens", "{max_new_tokens}", + "--output", + "results/research/k3_mac_bridge_vllm_mlx_niah.json", + ), + ), + timeout_minutes=60, + params={ + "n_samples": ("int:n_samples", "8"), + "max_new_tokens": ("int:max_new_tokens", "24"), + }, + validate_reports=False, + ), ) } diff --git a/scripts/research/vllm_mlx_niah_bench.py b/scripts/research/vllm_mlx_niah_bench.py new file mode 100644 index 00000000..dcd6cf1f --- /dev/null +++ b/scripts/research/vllm_mlx_niah_bench.py @@ -0,0 +1,354 @@ +#!/usr/bin/env python3 +"""vLLM-MLX parallel-NIAH probe (Mac bridge preset ``vllm-mlx-niah``). + +Answers one question on the SAME local MLX gemma verifier we used to reproduce +the MLX ``B>1, L=1`` batched-decode recall bug: **is vLLM-MLX both parallel AND +recall-preserving on our config?** + +It launches ``vllm-mlx serve --continuous-batching`` on the given model, fires +``--sessions`` concurrent needle-in-a-haystack requests (each with its OWN unique +needle, so a batching/cross-talk bug shows up as a recall drop), and reports: + + * per-session recall (needle found in that session's answer), + * aggregate decode tok/s at N concurrent vs an N=1 baseline (parallel speedup). + +Stdlib only (urllib + threads + subprocess); vLLM-MLX is the served process. The +script ALWAYS writes a verdict JSON (``status`` field) — even on install/load/ +server failure — so the bridge round-trip returns a usable answer, not a crash. +""" + +from __future__ import annotations + +import argparse +import json +import random +import socket +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + + +def _log(msg: str) -> None: + print(f"[vllm-mlx-niah] {msg}", file=sys.stderr, flush=True) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return int(s.getsockname()[1]) + + +def _vllm_mlx_version() -> Optional[str]: + try: + import importlib.metadata as m + + return m.version("vllm-mlx") + except Exception: + return None + + +def _build_niah_items(sessions: int, haystack_lines: int) -> List[Dict[str, str]]: + """One independent NIAH item per session, each with a UNIQUE access code.""" + rng = random.Random(1234) + cities = [ + "Lima", "Oslo", "Cairo", "Tokyo", "Quito", "Accra", "Riga", "Bern", + "Doha", "Suva", "Male", "Kyiv", "Vienna", "Hanoi", "Sofia", "Dakar", + ] + items: List[Dict[str, str]] = [] + for i in range(sessions): + code = f"{rng.randrange(16**6):06X}" # unique 6-hex code per session + filler = [ + f"On day {j}, the courier from {rng.choice(cities)} logged a " + f"routine delivery of crate {rng.randrange(1000)}." + for j in range(max(1, haystack_lines)) + ] + needle = f"IMPORTANT: the access code for vault {i} is {code}." + pos = rng.randrange(len(filler) + 1) + filler.insert(pos, needle) + prompt = ( + "Read the following log carefully.\n\n" + + "\n".join(filler) + + f"\n\nQuestion: What is the access code for vault {i}? " + "Answer with ONLY the code." + ) + items.append({"prompt": prompt, "code": code, "session": i}) + return items + + +def _get_model_id(base_url: str) -> Optional[str]: + """Resolve the served model id from /v1/models (the MLX verifier path).""" + try: + with urllib.request.urlopen(base_url + "/v1/models", timeout=10) as r: + data = json.loads(r.read().decode("utf-8")) + models = data.get("data") or [] + if models and isinstance(models[0], dict) and models[0].get("id"): + return str(models[0]["id"]) + except Exception: + pass + return None + + +def _extract(payload: Dict[str, Any]) -> Tuple[str, int]: + """Pull text + completion_tokens from a completions OR chat-completions body.""" + text = "" + try: + ch = payload["choices"][0] + text = ch.get("text") or (ch.get("message") or {}).get("content") or "" + except Exception: + text = "" + ctoks = 0 + usage = payload.get("usage") or {} + if isinstance(usage, dict) and isinstance(usage.get("completion_tokens"), int): + ctoks = usage["completion_tokens"] + if ctoks <= 0: + ctoks = max(1, len((text or "").split())) + return text or "", ctoks + + +def _post(base_url: str, path: str, body: Dict[str, Any], timeout: float): + req = urllib.request.Request( + base_url + path, data=json.dumps(body).encode("utf-8"), + headers={"Content-Type": "application/json"}, method="POST", + ) + with urllib.request.urlopen(req, timeout=timeout) as resp: + return json.loads(resp.read().decode("utf-8")) + + +def _post_generate( + base_url: str, model_id: str, prompt: str, max_new_tokens: int, timeout: float, +) -> Tuple[bool, str, int, float, str, str]: + """Generate via /v1/completions (raw prompt), falling back to chat. + + Returns (ok, text, completion_tokens, latency, err, endpoint). The MLX + verifier is a raw checkpoint (no chat template) so /v1/completions is the + primary path; chat is a fallback for instruct builds. + """ + t0 = time.time() + # gemma-4-IT needs its chat template; a raw /v1/completions prompt makes the + # instruct model emit immediately (empty answer). So try chat + # first (server applies the template), then fall back to /v1/completions with + # the gemma turn markers wrapped manually. + gemma = ( + f"user\n{prompt}\nmodel\n" + ) + attempts = ( + ("/v1/chat/completions", { + "model": model_id, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": int(max_new_tokens), "temperature": 0.0, + }), + ("/v1/completions", { + "model": model_id, "prompt": gemma, + "max_tokens": int(max_new_tokens), "temperature": 0.0, + "stop": [""], + }), + ) + last_err = "" + for path, body in attempts: + try: + payload = _post(base_url, path, body, timeout) + except urllib.error.HTTPError as exc: + last_err = f"{path} HTTP {exc.code}: {exc.read().decode('utf-8','replace')[:200]}" + continue # try the next endpoint shape + except Exception as exc: # noqa: BLE001 + return False, "", 0, time.time() - t0, f"{path}: {type(exc).__name__}: {exc}", path + text, ctoks = _extract(payload) + return True, text, ctoks, time.time() - t0, "", path + return False, "", 0, time.time() - t0, last_err, "none" + + +def _wait_for_server( + base_url: str, proc: subprocess.Popen, timeout: float, +) -> Tuple[bool, str]: + """Poll until the server answers, or it exits, or we time out.""" + deadline = time.time() + timeout + last = "" + while time.time() < deadline: + if proc.poll() is not None: + return False, f"server process exited early (rc={proc.returncode})" + for path in ("/v1/models", "/version", "/health"): + try: + with urllib.request.urlopen(base_url + path, timeout=5) as r: + if r.status == 200: + return True, path + except Exception as exc: # noqa: BLE001 + last = f"{type(exc).__name__}: {exc}" + time.sleep(3) + return False, f"timeout after {timeout:.0f}s (last: {last})" + + +def _tail(path: Path, n: int = 40) -> str: + try: + return "\n".join(path.read_text("utf-8", "replace").splitlines()[-n:]) + except Exception: + return "" + + +def _serve(model_path: str, port: int, continuous: bool, log_path: Path) -> subprocess.Popen: + argv = ["vllm-mlx", "serve", model_path, "--host", "127.0.0.1", + "--port", str(port), "--max-request-tokens", "32768"] + if continuous: + argv += ["--continuous-batching", "--use-paged-cache"] + _log(("continuous" if continuous else "simple") + " serve: " + " ".join(argv)) + with open(log_path, "wb") as logf: + return subprocess.Popen(argv, stdout=logf, stderr=subprocess.STDOUT) + + +def _run_phase( + model_path: str, continuous: bool, sessions: int, haystack_lines: int, + max_new_tokens: int, server_timeout: float, req_timeout: float, +) -> Dict[str, Any]: + """Launch one vLLM-MLX server (simple or continuous-batching) and run a + sessions-way concurrent NIAH against it. Returns a result dict.""" + phase: Dict[str, Any] = { + "continuous_batching": continuous, "sessions": sessions, "status": "init", + } + port = _free_port() + base_url = f"http://127.0.0.1:{port}" + log_path = Path(tempfile.gettempdir()) / f"vllm_mlx_serve_{port}.log" + proc: Optional[subprocess.Popen] = None + try: + proc = _serve(model_path, port, continuous, log_path) + ready, detail = _wait_for_server(base_url, proc, server_timeout) + if not ready: + phase.update(status="server_failed", error=detail, + server_log_tail=_tail(log_path)) + return phase + model_id = _get_model_id(base_url) or "default" + items = _build_niah_items(sessions, haystack_lines) + _post_generate(base_url, model_id, "Reply with the word ready.", 8, req_timeout) + + results: List[Any] = [None] * len(items) + t0 = time.time() + with ThreadPoolExecutor(max_workers=max(1, len(items))) as ex: + futs = {ex.submit(_post_generate, base_url, model_id, it["prompt"], + max_new_tokens, req_timeout): k + for k, it in enumerate(items)} + for fut in futs: + k = futs[fut] + try: + results[k] = fut.result() + except Exception as exc: # noqa: BLE001 + results[k] = (False, "", 0, 0.0, f"{type(exc).__name__}: {exc}", "none") + wall = max(time.time() - t0, 1e-6) + + per_session, hits, total_ctoks, n_ok = [], 0, 0, 0 + endpoint = "none" + for it, res in zip(items, results): + ok, text, ctoks, lat, err, ep = res + endpoint = ep if ep != "none" else endpoint + found = ok and (it["code"] in (text or "")) + hits += 1 if found else 0 + total_ctoks += ctoks if ok else 0 + n_ok += 1 if ok else 0 + per_session.append({ + "session": it["session"], "ok": ok, "needle_found": found, + "expected_code": it["code"], "answer_excerpt": (text or "")[:80], + "completion_tokens": ctoks, "latency_s": round(lat, 3), "error": err, + }) + phase.update( + status="ok", endpoint_used=endpoint, + recall=round(hits / len(items), 4) if items else 0.0, + sessions_ok=n_ok, total_completion_tokens=total_ctoks, + aggregate_decode_tps=round(total_ctoks / wall, 2), + wall_s=round(wall, 3), per_session=per_session, + server_log_tail=_tail(log_path, 12), + ) + return phase + except Exception as exc: # noqa: BLE001 + phase.update(status="error", error=f"{type(exc).__name__}: {exc}", + server_log_tail=_tail(log_path)) + return phase + finally: + if proc is not None and proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=20) + except Exception: + proc.kill() + time.sleep(2) # let the port/Metal context release before the next phase + + +def main() -> int: + ap = argparse.ArgumentParser(description="vLLM-MLX parallel NIAH probe") + ap.add_argument("--model-path", required=True, + help="Local MLX model dir (the same gemma verifier as the bug repro).") + ap.add_argument("--sessions", type=int, default=8) + ap.add_argument("--haystack-lines", type=int, default=60) + ap.add_argument("--max-new-tokens", type=int, default=24) + ap.add_argument("--server-timeout", type=float, default=900.0, + help="Seconds to wait for model load + server readiness.") + ap.add_argument("--req-timeout", type=float, default=300.0) + ap.add_argument("--output", required=True) + args = ap.parse_args() + + out = Path(args.output) + out.parent.mkdir(parents=True, exist_ok=True) + report: Dict[str, Any] = { + "kind": "vllm_mlx_niah_parallel", + "schema_version": 1, + "status": "init", + "config": { + "model_path": args.model_path, + "sessions": args.sessions, + "haystack_lines": args.haystack_lines, + "max_new_tokens": args.max_new_tokens, + "engine": "vllm-mlx serve --continuous-batching --use-paged-cache", + }, + "vllm_mlx_version": _vllm_mlx_version(), + } + + def _flush(status: str, **extra: Any) -> None: + report["status"] = status + report.update(extra) + out.write_text(json.dumps(report, indent=2), encoding="utf-8") + _log(f"status={status}; wrote {out}") + + if report["vllm_mlx_version"] is None: + _flush("vllm_mlx_not_installed", + error="`import importlib.metadata; version('vllm-mlx')` failed — " + "the pip-install step must run before this script.") + return 0 + _log(f"vllm-mlx version: {report['vllm_mlx_version']}") + + # Phase A — SIMPLE mode (no continuous batching): single-stream control. If + # gemma-4 generates here but fails under batching, the failure is isolated to + # vLLM-MLX's continuous-batching adapter, not model loading. + _log("=== Phase A: simple mode (single-stream control, N=1) ===") + simple = _run_phase( + args.model_path, continuous=False, sessions=1, + haystack_lines=args.haystack_lines, max_new_tokens=args.max_new_tokens, + server_timeout=args.server_timeout, req_timeout=args.req_timeout) + report["simple_mode"] = simple + + # Phase B — CONTINUOUS BATCHING (the parallel path under test): N sessions. + _log(f"=== Phase B: continuous batching, N={args.sessions} ===") + batched = _run_phase( + args.model_path, continuous=True, sessions=args.sessions, + haystack_lines=args.haystack_lines, max_new_tokens=args.max_new_tokens, + server_timeout=args.server_timeout, req_timeout=args.req_timeout) + report["continuous_batching_mode"] = batched + + # Verdict: parallel AND recall-preserving on our config? + simple_gen = simple.get("status") == "ok" and simple.get("total_completion_tokens", 0) > 0 + batched_recall = batched.get("recall", 0.0) if batched.get("status") == "ok" else 0.0 + batched_gen = batched.get("status") == "ok" and batched.get("total_completion_tokens", 0) > 0 + report["verdict"] = { + "single_stream_generates": bool(simple_gen), + "batched_generates": bool(batched_gen), + "batched_recall": batched_recall, + "parallel_and_recall_preserving": bool(batched_gen and batched_recall >= 0.99), + } + _flush("ok") + _log(f"VERDICT: {json.dumps(report['verdict'])}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/inference_engine/bridge/test_manifest.py b/tests/inference_engine/bridge/test_manifest.py index a287d708..e2e04888 100644 --- a/tests/inference_engine/bridge/test_manifest.py +++ b/tests/inference_engine/bridge/test_manifest.py @@ -84,6 +84,7 @@ def test_allowlist_contains_exactly_the_documented_presets(): "mlx-upgrade", "mlx-upstream-batch-probe", "pytest-path", + "vllm-mlx-niah", ] @@ -135,6 +136,25 @@ def test_pad_decode_preset_carries_flag_and_forces_trimmable_cache(): assert HARNESS_ENV["KAKEYA_MAC_VERIFIER_PATH"] in argv +def test_vllm_mlx_niah_preset_installs_then_benches(): + request = parse_manifest(_manifest( + preset="vllm-mlx-niah", params={"n_samples": "8", "max_new_tokens": "24"})) + commands = build_commands(request, HARNESS_ENV) + # two steps: pip install vllm-mlx, then the NIAH bench + assert len(commands) == 2 + pip = commands[0] + assert pip[:5] == ["python3", "-m", "pip", "install", "--upgrade"] + assert "vllm-mlx" in pip + bench = commands[1] + assert bench[1].endswith("vllm_mlx_niah_bench.py") + assert HARNESS_ENV["KAKEYA_MAC_VERIFIER_PATH"] in bench + assert bench[bench.index("--sessions") + 1] == "8" + assert bench[bench.index("--max-new-tokens") + 1] == "24" + # no unresolved placeholders survive + assert not [t for t in bench if t.startswith("${ENV:")] + assert not [t for t in bench if t.startswith("{") and t.endswith("}")] + + def test_drafter_parity_preset_resolves(): request = parse_manifest(_manifest( preset="k3-drafter-parity", params={"block_size": "8"}))