From e59de206806d7646d83cb19399e3ec24411678b3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 16 Jun 2026 11:05:28 +0000 Subject: [PATCH 1/5] =?UTF-8?q?feat(mac-bridge):=20vllm-mlx-niah=20preset?= =?UTF-8?q?=20=E2=80=94=20evaluate=20vLLM-MLX=20parallel=20+=20recall=20on?= =?UTF-8?q?=20Apple=20Silicon?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a Mac-bridge preset that runs vLLM-MLX (Apple-Silicon continuous batching) on the SAME local MLX gemma verifier used to reproduce the MLX B>1,L=1 batched- decode recall bug, to answer: is vLLM-MLX BOTH parallel AND recall-preserving on our config? (the Mac analog of CUDA KIE-v2 — Kakeya Attention on a borrowed runtime). - scripts/research/vllm_mlx_niah_bench.py: stdlib-only harness — installs/serves vllm-mlx --continuous-batching, fires N concurrent NIAH requests (unique needle per session so cross-talk shows as recall drop), reports per-session recall + aggregate decode tok/s vs N=1. Always writes a verdict JSON (status field). - manifest.py: vllm-mlx-niah preset (pip install vllm-mlx, then the bench); bounded params n_samples<=50, max_new_tokens<=512; model from runner env. - test_manifest.py: allowlist + build_commands coverage (100%, 28 tests). Co-authored-by: FluffyAIcode --- inference_engine/bridge/manifest.py | 30 ++ scripts/research/vllm_mlx_niah_bench.py | 292 ++++++++++++++++++ .../inference_engine/bridge/test_manifest.py | 20 ++ 3 files changed, 342 insertions(+) create mode 100644 scripts/research/vllm_mlx_niah_bench.py 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..c94dffbd --- /dev/null +++ b/scripts/research/vllm_mlx_niah_bench.py @@ -0,0 +1,292 @@ +#!/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 _post_chat( + base_url: str, prompt: str, max_new_tokens: int, timeout: float, +) -> Tuple[bool, str, int, float, str]: + """POST /v1/chat/completions. Returns (ok, text, completion_tokens, latency, err).""" + body = json.dumps({ + "model": "default", + "messages": [{"role": "user", "content": prompt}], + "max_tokens": int(max_new_tokens), + "temperature": 0.0, + }).encode("utf-8") + req = urllib.request.Request( + f"{base_url}/v1/chat/completions", data=body, + headers={"Content-Type": "application/json"}, method="POST", + ) + t0 = time.time() + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + payload = json.loads(resp.read().decode("utf-8")) + except urllib.error.HTTPError as exc: + detail = exc.read().decode("utf-8", "replace")[:300] + return False, "", 0, time.time() - t0, f"HTTP {exc.code}: {detail}" + except Exception as exc: # noqa: BLE001 - report any transport error + return False, "", 0, time.time() - t0, f"{type(exc).__name__}: {exc}" + latency = time.time() - t0 + try: + text = payload["choices"][0]["message"]["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: # fallback estimate if server omits usage + ctoks = max(1, len(text.split())) + return True, text, ctoks, latency, "" + + +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 ("/version", "/v1/models", "/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 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']}") + + port = _free_port() + base_url = f"http://127.0.0.1:{port}" + log_path = Path(tempfile.gettempdir()) / f"vllm_mlx_serve_{port}.log" + serve_argv = [ + "vllm-mlx", "serve", args.model_path, + "--host", "127.0.0.1", "--port", str(port), + "--continuous-batching", "--use-paged-cache", + "--max-request-tokens", "32768", + ] + _log("launching: " + " ".join(serve_argv)) + proc: Optional[subprocess.Popen] = None + try: + with open(log_path, "wb") as logf: + proc = subprocess.Popen(serve_argv, stdout=logf, stderr=subprocess.STDOUT) + + ready, detail = _wait_for_server(base_url, proc, args.server_timeout) + if not ready: + _flush("server_failed", + error=f"server not ready: {detail}", + server_log_tail=_tail(log_path)) + return 0 + _log(f"server ready ({detail})") + + items = _build_niah_items(args.sessions, args.haystack_lines) + + # Warmup (lazy MLX graph compile) — not measured. + _post_chat(base_url, "Reply with the word ready.", 8, args.req_timeout) + + # N=1 baseline (single request decode tok/s). + ok0, text0, ct0, lat0, err0 = _post_chat( + base_url, items[0]["prompt"], args.max_new_tokens, args.req_timeout) + n1_tps = (ct0 / lat0) if (ok0 and lat0 > 0) else 0.0 + + # N concurrent (the parallel path — continuous batching). + results: List[Optional[Tuple[bool, str, int, float, str]]] = [None] * len(items) + t0 = time.time() + with ThreadPoolExecutor(max_workers=len(items)) as ex: + futs = { + ex.submit(_post_chat, base_url, it["prompt"], + args.max_new_tokens, args.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}") + wall = max(time.time() - t0, 1e-6) + + per_session: List[Dict[str, Any]] = [] + hits = 0 + total_ctoks = 0 + n_ok = 0 + for k, (it, res) in enumerate(zip(items, results)): + ok, text, ctoks, lat, err = res # type: ignore[misc] + 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, + }) + + recall = hits / len(items) if items else 0.0 + agg_tps = total_ctoks / wall + _flush( + "ok", + recall=round(recall, 4), + sessions_ok=n_ok, + n1_decode_tps=round(n1_tps, 2), + aggregate_decode_tps=round(agg_tps, 2), + parallel_speedup_vs_n1=round(agg_tps / n1_tps, 3) if n1_tps > 0 else None, + concurrent_wall_s=round(wall, 3), + total_completion_tokens=total_ctoks, + per_session=per_session, + server_log_tail=_tail(log_path, 20), + ) + verdict = ( + f"recall={recall:.3f} ({hits}/{len(items)}), " + f"agg_decode={agg_tps:.1f} tok/s, N=1={n1_tps:.1f} tok/s, " + f"parallel={'YES' if agg_tps > n1_tps else 'no-gain'}" + ) + _log("VERDICT: " + verdict) + return 0 + except Exception as exc: # noqa: BLE001 + _flush("error", error=f"{type(exc).__name__}: {exc}", + server_log_tail=_tail(log_path)) + return 0 + finally: + if proc is not None and proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=20) + except Exception: + proc.kill() + + +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"})) From 1fa1ab85d901fa7d8097a17ccd2d4a690a7f5031 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 16 Jun 2026 11:13:02 +0000 Subject: [PATCH 2/5] fix(vllm-mlx-niah): use /v1/completions + resolved model id (chat 404 on raw MLX verifier) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First Mac run: vllm-mlx 0.3.0 installed + loaded gemma-4-26B-A4B-it-mlx-4bit and served (continuous batching), but every /v1/chat/completions 404'd (the MLX verifier is a raw checkpoint with no chat template; /version also 404). Harness now resolves the served model id from /v1/models and generates via /v1/completions (raw prompt), falling back to chat — so recall/throughput measure the real parallel path. Co-authored-by: FluffyAIcode --- scripts/research/vllm_mlx_niah_bench.py | 116 +++++++++++++++++------- 1 file changed, 81 insertions(+), 35 deletions(-) diff --git a/scripts/research/vllm_mlx_niah_bench.py b/scripts/research/vllm_mlx_niah_bench.py index c94dffbd..12ac5a0d 100644 --- a/scripts/research/vllm_mlx_niah_bench.py +++ b/scripts/research/vllm_mlx_niah_bench.py @@ -81,41 +81,78 @@ def _build_niah_items(sessions: int, haystack_lines: int) -> List[Dict[str, str] return items -def _post_chat( - base_url: str, prompt: str, max_new_tokens: int, timeout: float, -) -> Tuple[bool, str, int, float, str]: - """POST /v1/chat/completions. Returns (ok, text, completion_tokens, latency, err).""" - body = json.dumps({ - "model": "default", - "messages": [{"role": "user", "content": prompt}], - "max_tokens": int(max_new_tokens), - "temperature": 0.0, - }).encode("utf-8") - req = urllib.request.Request( - f"{base_url}/v1/chat/completions", data=body, - headers={"Content-Type": "application/json"}, method="POST", - ) - t0 = time.time() +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(req, timeout=timeout) as resp: - payload = json.loads(resp.read().decode("utf-8")) - except urllib.error.HTTPError as exc: - detail = exc.read().decode("utf-8", "replace")[:300] - return False, "", 0, time.time() - t0, f"HTTP {exc.code}: {detail}" - except Exception as exc: # noqa: BLE001 - report any transport error - return False, "", 0, time.time() - t0, f"{type(exc).__name__}: {exc}" - latency = time.time() - t0 + 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: - text = payload["choices"][0]["message"]["content"] or "" + 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: # fallback estimate if server omits usage - ctoks = max(1, len(text.split())) - return True, text, ctoks, latency, "" + 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() + attempts = ( + ("/v1/completions", { + "model": model_id, "prompt": prompt, + "max_tokens": int(max_new_tokens), "temperature": 0.0, + }), + ("/v1/chat/completions", { + "model": model_id, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": int(max_new_tokens), "temperature": 0.0, + }), + ) + 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( @@ -127,7 +164,7 @@ def _wait_for_server( while time.time() < deadline: if proc.poll() is not None: return False, f"server process exited early (rc={proc.returncode})" - for path in ("/version", "/v1/models", "/health"): + for path in ("/v1/models", "/version", "/health"): try: with urllib.request.urlopen(base_url + path, timeout=5) as r: if r.status == 200: @@ -210,22 +247,31 @@ def _flush(status: str, **extra: Any) -> None: return 0 _log(f"server ready ({detail})") + model_id = _get_model_id(base_url) or "default" + report["served_model_id"] = model_id + _log(f"served model id: {model_id}") + items = _build_niah_items(args.sessions, args.haystack_lines) # Warmup (lazy MLX graph compile) — not measured. - _post_chat(base_url, "Reply with the word ready.", 8, args.req_timeout) + _post_generate(base_url, model_id, "Reply with the word ready.", + 8, args.req_timeout) # N=1 baseline (single request decode tok/s). - ok0, text0, ct0, lat0, err0 = _post_chat( - base_url, items[0]["prompt"], args.max_new_tokens, args.req_timeout) + ok0, text0, ct0, lat0, err0, ep0 = _post_generate( + base_url, model_id, items[0]["prompt"], args.max_new_tokens, + args.req_timeout) n1_tps = (ct0 / lat0) if (ok0 and lat0 > 0) else 0.0 + report["endpoint_used"] = ep0 + if not ok0: + report["n1_error"] = err0 # N concurrent (the parallel path — continuous batching). - results: List[Optional[Tuple[bool, str, int, float, str]]] = [None] * len(items) + results: List[Optional[Tuple[bool, str, int, float, str, str]]] = [None] * len(items) t0 = time.time() with ThreadPoolExecutor(max_workers=len(items)) as ex: futs = { - ex.submit(_post_chat, base_url, it["prompt"], + ex.submit(_post_generate, base_url, model_id, it["prompt"], args.max_new_tokens, args.req_timeout): k for k, it in enumerate(items) } @@ -234,7 +280,7 @@ def _flush(status: str, **extra: Any) -> None: try: results[k] = fut.result() except Exception as exc: # noqa: BLE001 - results[k] = (False, "", 0, 0.0, f"{type(exc).__name__}: {exc}") + results[k] = (False, "", 0, 0.0, f"{type(exc).__name__}: {exc}", "none") wall = max(time.time() - t0, 1e-6) per_session: List[Dict[str, Any]] = [] @@ -242,7 +288,7 @@ def _flush(status: str, **extra: Any) -> None: total_ctoks = 0 n_ok = 0 for k, (it, res) in enumerate(zip(items, results)): - ok, text, ctoks, lat, err = res # type: ignore[misc] + ok, text, ctoks, lat, err, _ep = res # type: ignore[misc] found = ok and (it["code"] in (text or "")) hits += 1 if found else 0 total_ctoks += ctoks if ok else 0 From 7678e4cc987fcc4c6d54a7528347a89d9b1d2973 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 16 Jun 2026 11:17:50 +0000 Subject: [PATCH 3/5] =?UTF-8?q?fix(vllm-mlx-niah):=20gemma-4-IT=20needs=20?= =?UTF-8?q?chat=20template=20=E2=80=94=20try=20chat=20first,=20wrap=20comp?= =?UTF-8?q?letions=20fallback=20in=20turn=20markers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Run 2: /v1/completions returned empty (1 token) — gemma-4-it emits immediately on a raw prompt. Now try /v1/chat/completions first (server applies the template) with the resolved model id, falling back to /v1/completions with the prompt wrapped in ...model markers + an stop. Batching already confirmed active (8/8 ok, 5.06x scaling). Co-authored-by: FluffyAIcode --- scripts/research/vllm_mlx_niah_bench.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/scripts/research/vllm_mlx_niah_bench.py b/scripts/research/vllm_mlx_niah_bench.py index 12ac5a0d..c7580a7f 100644 --- a/scripts/research/vllm_mlx_niah_bench.py +++ b/scripts/research/vllm_mlx_niah_bench.py @@ -130,16 +130,24 @@ def _post_generate( 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/completions", { - "model": model_id, "prompt": prompt, - "max_tokens": int(max_new_tokens), "temperature": 0.0, - }), ("/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: From ace720fddb4477673e382f27e16a56ee56cbde2f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 16 Jun 2026 11:22:46 +0000 Subject: [PATCH 4/5] debug(vllm-mlx-niah): capture raw server response for a control prompt (diagnose empty answer) Co-authored-by: FluffyAIcode --- scripts/research/vllm_mlx_niah_bench.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/scripts/research/vllm_mlx_niah_bench.py b/scripts/research/vllm_mlx_niah_bench.py index c7580a7f..aaa9392f 100644 --- a/scripts/research/vllm_mlx_niah_bench.py +++ b/scripts/research/vllm_mlx_niah_bench.py @@ -259,6 +259,30 @@ def _flush(status: str, **extra: Any) -> None: report["served_model_id"] = model_id _log(f"served model id: {model_id}") + # DEBUG: capture the raw server response for a simple control prompt, so + # we can see finish_reason / structure (diagnose the empty-answer issue). + for dbgname, dbgbody in ( + ("debug_simple_chat", { + "model": model_id, + "messages": [{"role": "user", + "content": "What is the capital of France? Answer in one short sentence."}], + "max_tokens": 32, "temperature": 0.0}), + ("debug_simple_completion", { + "model": model_id, + "prompt": "user\nWhat is the capital of France?\nmodel\n", + "max_tokens": 32, "temperature": 0.0, "stop": [""]}), + ): + path = ("/v1/chat/completions" if "chat" in dbgname + else "/v1/completions") + try: + dbg = _post(base_url, path, dbgbody, args.req_timeout) + report[dbgname] = json.dumps(dbg)[:900] + except urllib.error.HTTPError as exc: + report[dbgname] = f"HTTP {exc.code}: {exc.read().decode('utf-8','replace')[:300]}" + except Exception as exc: # noqa: BLE001 + report[dbgname] = f"{type(exc).__name__}: {exc}" + _log(f"{dbgname}: {report[dbgname][:300]}") + items = _build_niah_items(args.sessions, args.haystack_lines) # Warmup (lazy MLX graph compile) — not measured. From b6afefa54ac01ee8a3c7880304ea1341c9cd88fd Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 16 Jun 2026 11:29:15 +0000 Subject: [PATCH 5/5] =?UTF-8?q?feat(vllm-mlx-niah):=20two-phase=20A/B=20?= =?UTF-8?q?=E2=80=94=20simple=20single-stream=20control=20vs=20continuous-?= =?UTF-8?q?batching=20N=3D8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Run 4 evidence: vllm-mlx 0.3.0 errors under continuous batching on gemma-4 (patch_gemma4_attention_for_batching got unexpected kwarg 'shared_kv' → finish_reason=error, 0 tokens). Now run BOTH phases in one go: Phase A simple mode (no continuous batching, N=1) isolates whether gemma-4 generates at all, Phase B continuous batching N=8 is the parallel+recall test. Emits a verdict: single_stream_generates / batched_generates / batched_recall / parallel_and_recall_preserving. Co-authored-by: FluffyAIcode --- scripts/research/vllm_mlx_niah_bench.py | 248 +++++++++++------------- 1 file changed, 116 insertions(+), 132 deletions(-) diff --git a/scripts/research/vllm_mlx_niah_bench.py b/scripts/research/vllm_mlx_niah_bench.py index aaa9392f..dcd6cf1f 100644 --- a/scripts/research/vllm_mlx_niah_bench.py +++ b/scripts/research/vllm_mlx_niah_bench.py @@ -190,6 +190,91 @@ def _tail(path: Path, n: int = 40) -> str: 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, @@ -232,138 +317,37 @@ def _flush(status: str, **extra: Any) -> None: return 0 _log(f"vllm-mlx version: {report['vllm_mlx_version']}") - port = _free_port() - base_url = f"http://127.0.0.1:{port}" - log_path = Path(tempfile.gettempdir()) / f"vllm_mlx_serve_{port}.log" - serve_argv = [ - "vllm-mlx", "serve", args.model_path, - "--host", "127.0.0.1", "--port", str(port), - "--continuous-batching", "--use-paged-cache", - "--max-request-tokens", "32768", - ] - _log("launching: " + " ".join(serve_argv)) - proc: Optional[subprocess.Popen] = None - try: - with open(log_path, "wb") as logf: - proc = subprocess.Popen(serve_argv, stdout=logf, stderr=subprocess.STDOUT) - - ready, detail = _wait_for_server(base_url, proc, args.server_timeout) - if not ready: - _flush("server_failed", - error=f"server not ready: {detail}", - server_log_tail=_tail(log_path)) - return 0 - _log(f"server ready ({detail})") - - model_id = _get_model_id(base_url) or "default" - report["served_model_id"] = model_id - _log(f"served model id: {model_id}") - - # DEBUG: capture the raw server response for a simple control prompt, so - # we can see finish_reason / structure (diagnose the empty-answer issue). - for dbgname, dbgbody in ( - ("debug_simple_chat", { - "model": model_id, - "messages": [{"role": "user", - "content": "What is the capital of France? Answer in one short sentence."}], - "max_tokens": 32, "temperature": 0.0}), - ("debug_simple_completion", { - "model": model_id, - "prompt": "user\nWhat is the capital of France?\nmodel\n", - "max_tokens": 32, "temperature": 0.0, "stop": [""]}), - ): - path = ("/v1/chat/completions" if "chat" in dbgname - else "/v1/completions") - try: - dbg = _post(base_url, path, dbgbody, args.req_timeout) - report[dbgname] = json.dumps(dbg)[:900] - except urllib.error.HTTPError as exc: - report[dbgname] = f"HTTP {exc.code}: {exc.read().decode('utf-8','replace')[:300]}" - except Exception as exc: # noqa: BLE001 - report[dbgname] = f"{type(exc).__name__}: {exc}" - _log(f"{dbgname}: {report[dbgname][:300]}") - - items = _build_niah_items(args.sessions, args.haystack_lines) - - # Warmup (lazy MLX graph compile) — not measured. - _post_generate(base_url, model_id, "Reply with the word ready.", - 8, args.req_timeout) - - # N=1 baseline (single request decode tok/s). - ok0, text0, ct0, lat0, err0, ep0 = _post_generate( - base_url, model_id, items[0]["prompt"], args.max_new_tokens, - args.req_timeout) - n1_tps = (ct0 / lat0) if (ok0 and lat0 > 0) else 0.0 - report["endpoint_used"] = ep0 - if not ok0: - report["n1_error"] = err0 - - # N concurrent (the parallel path — continuous batching). - results: List[Optional[Tuple[bool, str, int, float, str, str]]] = [None] * len(items) - t0 = time.time() - with ThreadPoolExecutor(max_workers=len(items)) as ex: - futs = { - ex.submit(_post_generate, base_url, model_id, it["prompt"], - args.max_new_tokens, args.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: List[Dict[str, Any]] = [] - hits = 0 - total_ctoks = 0 - n_ok = 0 - for k, (it, res) in enumerate(zip(items, results)): - ok, text, ctoks, lat, err, _ep = res # type: ignore[misc] - 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, - }) - - recall = hits / len(items) if items else 0.0 - agg_tps = total_ctoks / wall - _flush( - "ok", - recall=round(recall, 4), - sessions_ok=n_ok, - n1_decode_tps=round(n1_tps, 2), - aggregate_decode_tps=round(agg_tps, 2), - parallel_speedup_vs_n1=round(agg_tps / n1_tps, 3) if n1_tps > 0 else None, - concurrent_wall_s=round(wall, 3), - total_completion_tokens=total_ctoks, - per_session=per_session, - server_log_tail=_tail(log_path, 20), - ) - verdict = ( - f"recall={recall:.3f} ({hits}/{len(items)}), " - f"agg_decode={agg_tps:.1f} tok/s, N=1={n1_tps:.1f} tok/s, " - f"parallel={'YES' if agg_tps > n1_tps else 'no-gain'}" - ) - _log("VERDICT: " + verdict) - return 0 - except Exception as exc: # noqa: BLE001 - _flush("error", error=f"{type(exc).__name__}: {exc}", - server_log_tail=_tail(log_path)) - return 0 - finally: - if proc is not None and proc.poll() is None: - proc.terminate() - try: - proc.wait(timeout=20) - except Exception: - proc.kill() + # 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__":