From 350846634db67631016841aea96351e6de774514 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 22 Jul 2026 20:24:52 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20Xe2=20performance=20layer=20=E2=80=94?= =?UTF-8?q?=20XMX=20flash-attention=20recipes,=20managed=20JIT=20cache,=20?= =?UTF-8?q?autotune?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four optimizations not available in any other Intel Arc runtime: - XMX flash-attention aware recipes: llama.cpp's oneDNN SDPA path (merged 2026-07-15, up to 4.26x prefill at 80k ctx on Battlemage) requires an f16 KV cache. On Xe2, default recipes now prefer f16 KV + -fa on + -ub 1024 -b 2048 whenever that still leaves >=16k ctx, falling back to q8_0 KV on VRAM-tight setups. New LaunchRecipe fields: flash_attn, batch_size, cache_reuse, no_mmap. - Managed persistent SYCL JIT cache (sycl_cache.py): replaces the blanket SYCL_CACHE_PERSISTENT=0 workaround (~20s JIT per cold start) with a per-(binary, driver) fingerprinted cache dir plus a warm-up crash guard — a run that dies before its first health check wipes and poisons that cache and falls back to the old behaviour. Opt out: [server] jit_cache="off". - arc-llama tune: on-device autotune sweeping q8-KV vs f16+FA across ubatch sizes through the live server, scored by modelled request latency (prefill + generation), winner persisted to the model recipe. The /admin/models/{name}/edit endpoint now accepts batch_size, flash_attn, cache_reuse, and null values to clear optional fields. - Slow-quant warning: registering a Q8_0 weight quant on Xe2 warns about the known ~4-5x kernel-efficiency gap vs Q4_K_M (ggml-org/llama.cpp#21517). Non-SWA recipes also gain --cache-reuse 256 for multi-turn TTFT. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Ezf1mE6Cd15wgCFMCTBZrM --- README.md | 46 ++++++++- src/arc_llama/arch.py | 31 +++++- src/arc_llama/benchmark.py | 183 ++++++++++++++++++++++++++++++++++++ src/arc_llama/cli.py | 67 +++++++++++++ src/arc_llama/config.py | 10 ++ src/arc_llama/launcher.py | 59 +++++++++++- src/arc_llama/models.py | 73 +++++++++++--- src/arc_llama/recipes.py | 73 +++++++++++++- src/arc_llama/server.py | 56 +++++++++-- src/arc_llama/sycl_cache.py | 181 +++++++++++++++++++++++++++++++++++ tests/test_launcher.py | 48 ++++++++++ tests/test_models.py | 41 ++++++++ tests/test_recipes.py | 84 ++++++++++++++++- tests/test_sycl_cache.py | 107 +++++++++++++++++++++ 14 files changed, 1028 insertions(+), 31 deletions(-) create mode 100644 src/arc_llama/sycl_cache.py create mode 100644 tests/test_sycl_cache.py diff --git a/README.md b/README.md index 1885845..7c83766 100644 --- a/README.md +++ b/README.md @@ -28,8 +28,9 @@ something useful before lunch. - **Auto-discovery** of every Intel GPU on the host (`Alchemist`, `Battlemage`, Lunar Lake iGPU). PCI device-ID table covers the common SKUs and falls back to OpenCL device-name parsing for the rest. -- **Per-arch SYCL profiles** , env vars like `SYCL_CACHE_PERSISTENT=0` are - applied automatically, and known-bad ones (e.g. `GGML_SYCL_DISABLE_OPT`, +- **Per-arch SYCL profiles** , env vars like `SYCL_CACHE_PERSISTENT` are + managed automatically (see the managed JIT cache below), and known-bad ones + (e.g. `GGML_SYCL_DISABLE_OPT`, `SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS`) are stripped from the inherited shell environment. - **Smart defaults** for `-ctx`, `--cache-type-k/v`, and `-ngl` based on the @@ -109,6 +110,47 @@ curl http://127.0.0.1:11437/v1/chat/completions \ cmake --build build --config Release -j ``` - User in the `render` and `video` groups (`arc-llama doctor` will tell you). +- For the XMX flash-attention prefill path (below), add `-DGGML_SYCL_DNN=ON` + (oneDNN ≥ 3.11.2, llama.cpp ≥ b10016 / 2026-07-15). + +## Performance: beyond stock settings + +arc-llama doesn't just avoid the Arc footguns — it applies optimizations no +other Intel runtime ships yet: + +- **XMX flash-attention aware recipes.** llama.cpp's oneDNN SDPA path + (merged 2026-07-15) routes long prefills through the XMX systolic arrays — + up to **4.26× faster prefill at 80k ctx** on Battlemage — but only with an + **f16 KV cache**; quantized KV silently falls back to shader kernels. On + Xe2, arc-llama now defaults to f16 KV + `-fa on` + `-ub 1024 -b 2048` + whenever VRAM still allows ≥16k context, and drops back to q8_0 KV (more + context per byte) only on VRAM-tight setups. Every other tool leaves you + with the pre-XMX defaults — quantized KV that quietly disables the fastest + attention path Intel hardware has. +- **Managed persistent SYCL JIT cache.** Battlemage + oneAPI 2026.0 SIGSEGVs + when `SYCL_CACHE_PERSISTENT=1` reads stale cache entries; everyone works + around it by disabling the cache — paying **~20 s of JIT recompilation on + every cold start, forever**. arc-llama instead gives each + (llama-server build, GPU driver) combination its own fingerprinted cache + dir, so stale entries can't exist, and arms a crash guard: if a server ever + dies during warm-up with the cache active, that cache is wiped, poisoned, + and the safe `=0` behaviour returns automatically. Net effect: the JIT cost + is paid once per llama.cpp upgrade instead of on every model swap. + Opt out with `jit_cache = "off"` under `[server]`. +- **Multi-turn prompt reuse.** Non-SWA recipes get `--cache-reuse 256`, so + follow-up chat turns re-use the shared-prefix KV instead of re-prefilling + the whole conversation. +- **`arc-llama tune`** — on-device autotuning. With `arc-llama serve` + running, `arc-llama tune ` sweeps the configs that actually matter + on Arc (q8 KV vs f16+FA, ubatch 512/1024/2048), measures real prefill + + generation through your own stack, scores each by modelled request latency, + and persists the winner into the model's recipe. Your card, your driver, + your build — not someone else's benchmark table. +- **Slow-quant warnings.** Q8_0 *weight* quants hit a known-bad kernel path + on Xe2 (~22% of memory bandwidth vs Q4_K_M's ~55%, + [llama.cpp#21517](https://github.com/ggml-org/llama.cpp/issues/21517)) — + registering one warns you to grab Q4_K_M/Q6_K instead of leaving you to + wonder why generation crawls. ## Multi-GPU diff --git a/src/arc_llama/arch.py b/src/arc_llama/arch.py index bb6ccb3..07136b8 100644 --- a/src/arc_llama/arch.py +++ b/src/arc_llama/arch.py @@ -32,6 +32,19 @@ class ArchProfile: """Whether q8_0 K/V cache produces correct generation on this arch.""" prefer_uniform_quants: bool = True """If true, recommend Q4_K_M over Unsloth Dynamic XL/UD variants.""" + supports_xmx_fa: bool = False + """Whether the oneDNN XMX flash-attention prefill path (llama.cpp + PR #25222, merged 2026-07-15) is confirmed working on this arch. + The path requires an f16 K/V cache — quantized KV silently falls back + to the shader kernels, so recipes trade KV size for prefill speed.""" + slow_weight_quants: tuple[str, ...] = () + """Weight-quant tiers with known-bad kernel efficiency on this arch. + Used to warn at registration time, e.g. Q8_0 on Xe2 runs at ~22% of + memory bandwidth vs Q4_K_M's ~55% (ggml-org/llama.cpp#21517).""" + managed_jit_cache: bool = False + """If true, arc-llama replaces this profile's SYCL_CACHE_PERSISTENT=0 + workaround with a managed per-binary cache dir (see sycl_cache.py), + recovering the ~20s cold-start JIT cost without the stale-cache SIGSEGV.""" # --------------------------------------------------------------------------- @@ -150,12 +163,22 @@ class ArchProfile: notes=[ "Requires kernel 6.14+ and Mesa 24.x+ for stable `xe` driver.", "ReBAR REQUIRED — without it llama.cpp will fall back to slow paths.", - "First inference per cold start pays ~20s of SYCL JIT compile.", + "First inference per cold start pays ~20s of SYCL JIT compile " + "(recovered automatically when the managed JIT cache is active).", "q8_0 K/V cache works correctly but on some builds underutilises memory " "bandwidth on dense models. Verify perf if you care; correctness is OK.", + "llama.cpp builds with GGML_SYCL_DNN=ON route long prefills through the " + "XMX engines when flash attention is on and the KV cache is f16 — up to " + "4.26x prefill at 80k ctx (PR #25222). arc-llama recipes prefer f16 KV " + "on this arch when VRAM allows.", + "Avoid Q8_0 *weight* quants here — Q4_K_M/Q4_K_S/Q6_K run 4-5x faster " + "(llama.cpp issue #21517).", ], safe_kv_q8=True, prefer_uniform_quants=True, + supports_xmx_fa=True, + slow_weight_quants=("Q8_0",), + managed_jit_cache=True, ) LUNAR_LAKE_PROFILE = ArchProfile( @@ -177,6 +200,12 @@ class ArchProfile: ], safe_kv_q8=True, prefer_uniform_quants=True, + # Xe2-LPG carries the same XMX engines and JIT-cache behaviour as + # Battlemage; the shared-memory VRAM budget is what differs, and that is + # already handled by ctx sizing. + supports_xmx_fa=True, + slow_weight_quants=("Q8_0",), + managed_jit_cache=True, ) UNKNOWN_PROFILE = ArchProfile( diff --git a/src/arc_llama/benchmark.py b/src/arc_llama/benchmark.py index fae87e8..90af1dc 100644 --- a/src/arc_llama/benchmark.py +++ b/src/arc_llama/benchmark.py @@ -366,6 +366,189 @@ async def benchmark_sweep( return results +# ------------------------------------------------------------------ +# Autotune +# ------------------------------------------------------------------ + +TUNE_PROMPT_TOKENS = 2048 +TUNE_GEN_TOKENS = 128 +"""Autotune workload shape: long enough prefill to expose the XMX/oneDNN +flash-attention path (which only engages ≥32-token query batches and pays +off at scale), plus a chat-sized generation burst.""" + + +@dataclass +class TuneCandidate: + """One recipe variant the autotuner will measure.""" + label: str + edit: dict[str, Any] + + +@dataclass +class TuneOutcome: + candidate: TuneCandidate + result: BenchmarkResult + + @property + def request_seconds(self) -> float | None: + """Modelled wall time of one real request: prefill + generation. + + This is the scoring metric — it weighs prompt and generation speed by + how much time each actually costs, instead of an arbitrary blend. + """ + pp = self.result.prompt_eval_tok_s + gen = self.result.generation_tok_s + if not pp or not gen: + return None + return self.result.prompt_tokens / pp + self.result.gen_tokens / gen + + +def build_tune_candidates(arch_value: str, current_recipe: dict[str, Any]) -> list[TuneCandidate]: + """Candidate grid for one model on one arch. + + Kept deliberately small — each candidate costs a full model reload plus + JIT warm-up, so we probe the configurations that actually move the needle + on Arc rather than a blind grid: + * q8_0 KV without forced FA (the pre-XMX-era default, still best when + the quantized-KV shader path happens to win on a given build); + * f16 KV + -fa on across ubatch 512/1024/2048 (the oneDNN XMX prefill + path needs f16 KV; ubatch controls how well-fed the systolic arrays are). + """ + from arc_llama.arch import Arch, profile_for + try: + profile = profile_for(Arch(arch_value)) + except ValueError: + profile = profile_for(Arch.UNKNOWN) + + cands = [TuneCandidate( + label="current", + edit={ + k: current_recipe[k] + for k in ("cache_type_k", "cache_type_v", "ubatch_size", "batch_size", "flash_attn") + if k in current_recipe + } or {"flash_attn": None}, + )] + cands.append(TuneCandidate( + label="q8-kv", + edit={"cache_type_k": "q8_0", "cache_type_v": "q8_0", + "flash_attn": None, "ubatch_size": 512, "batch_size": None}, + )) + ubatches = [512, 1024, 2048] if profile.supports_xmx_fa else [512, 1024] + for ub in ubatches: + cands.append(TuneCandidate( + label=f"f16-fa-ub{ub}", + edit={"cache_type_k": "f16", "cache_type_v": "f16", + "flash_attn": "on", "ubatch_size": ub, "batch_size": max(2048, ub)}, + )) + return cands + + +async def autotune_model( + server_url: str, + model_name: str, + *, + prompt_tokens: int = TUNE_PROMPT_TOKENS, + gen_tokens: int = TUNE_GEN_TOKENS, + apply_best: bool = True, + cfg: Config | None = None, +) -> tuple[list[TuneOutcome], TuneOutcome | None]: + """Measure each candidate on the real hardware and keep the fastest. + + Every candidate goes through the running arc-llama server (correct SYCL + env, arch profile, managed JIT cache), gets a full reload, warm-up, and + prefill+generation measurement. Winner = lowest modelled request time. + Applies the winner to the model's recipe unless apply_best=False, in + which case the original recipe is restored. + """ + if cfg is None: + cfg = load_config() + model = cfg.find_model(model_name) + if model is None: + raise ValueError(f"Model '{model_name}' not found in config") + gpu = cfg.find_gpu(model.gpu_pci_slot) + arch_value = gpu.arch if gpu else "unknown" + original_recipe = dict(model.recipe or {}) + + candidates = build_tune_candidates(arch_value, original_recipe) + outcomes: list[TuneOutcome] = [] + # MTP models pin ubatch to 8 — sweeping prefill batches would fight the + # SSM compute-buffer constraint, so only the current config is measured. + if original_recipe.get("spec_type") == "draft-mtp": + log.warning("%s uses draft-mtp; skipping batch/FA sweep, keeping current recipe", model_name) + candidates = candidates[:1] + + async with httpx.AsyncClient(base_url=server_url, timeout=600.0) as client: + for cand in candidates: + log.info("tuning %s: %s ...", model_name, cand.label) + r = await client.post(f"/admin/models/{model_name}/edit", json=cand.edit) + if r.status_code != 200: + outcomes.append(TuneOutcome(cand, BenchmarkResult( + model=model_name, ctx=0, + cache_type_k=str(cand.edit.get("cache_type_k", "?")), + cache_type_v=str(cand.edit.get("cache_type_v", "?")), + prompt_tokens=prompt_tokens, gen_tokens=gen_tokens, + error=f"edit failed: {r.status_code} {r.text}", + ))) + continue + res = await benchmark_model( + server_url, model_name, + prompt_tokens=prompt_tokens, gen_tokens=gen_tokens, + load=True, cfg=cfg, + ) + outcomes.append(TuneOutcome(cand, res)) + + scored = [o for o in outcomes if o.request_seconds is not None and not o.result.error] + best = min(scored, key=lambda o: o.request_seconds) if scored else None + + # Land on the winner (or roll back). The edit endpoint persists to + # config.toml, so whatever we set last is what survives. + final_edit = best.candidate.edit if (best and apply_best) else { + k: original_recipe.get(k) + for k in ("cache_type_k", "cache_type_v", "ubatch_size", "batch_size", "flash_attn") + } + final_edit = {k: v for k, v in final_edit.items() if k in ( + "cache_type_k", "cache_type_v", "ubatch_size", "batch_size", "flash_attn", + )} + # cache_type must never be null; drop absent keys instead. + for k in ("cache_type_k", "cache_type_v"): + if final_edit.get(k) is None: + final_edit.pop(k, None) + if final_edit: + await client.post(f"/admin/models/{model_name}/edit", json=final_edit) + + return outcomes, best + + +def print_tune_table(outcomes: list[TuneOutcome], best: TuneOutcome | None) -> None: + from rich.console import Console + from rich.table import Table + + console = Console() + table = Table(title="Autotune") + table.add_column("candidate") + table.add_column("KV") + table.add_column("Prompt-eval") + table.add_column("Generation") + table.add_column("2k-req time") + table.add_column("") + for o in outcomes: + r = o.result + if r.error: + table.add_row(o.candidate.label, f"{r.cache_type_k}/{r.cache_type_v}", + "[red]error[/red]", "", "", "") + continue + secs = o.request_seconds + table.add_row( + o.candidate.label, + f"{r.cache_type_k}/{r.cache_type_v}", + _fmt_speed(r.prompt_eval_tok_s), + _fmt_speed(r.generation_tok_s), + f"{secs:.2f} s" if secs is not None else "—", + "◀ winner" if best is o else "", + ) + console.print(table) + + # ------------------------------------------------------------------ # Formatting # ------------------------------------------------------------------ diff --git a/src/arc_llama/cli.py b/src/arc_llama/cli.py index b24267c..49e6021 100644 --- a/src/arc_llama/cli.py +++ b/src/arc_llama/cli.py @@ -9,6 +9,7 @@ arc-llama add Register a model — local file or HF download. arc-llama remove Remove a model from the config. arc-llama serve Run the OpenAI-compatible router. + arc-llama tune Autotune a model's recipe on the actual hardware. arc-llama tui Launch the terminal UI. arc-llama systemd Print a systemd --user service unit for `arc-llama serve`. @@ -615,6 +616,72 @@ def _on_signal(signum: int, _frame) -> None: # noqa: ANN001 uvicorn.run(app, host=cfg.server.host, port=cfg.server.port, log_level="info") +# =========================================================================== +# tune +# =========================================================================== + +@cli.command("tune") +@click.argument("model") +@click.option("--server-url", default=None, + help="Base URL of the running `arc-llama serve` (default: from config).") +@click.option("--prompt-tokens", type=int, default=None, + help="Prefill size to benchmark (default: 2048).") +@click.option("--gen-tokens", type=int, default=None, + help="Generation burst to benchmark (default: 128).") +@click.option("--no-apply", is_flag=True, + help="Measure only; restore the original recipe afterwards.") +@click.pass_context +def tune_cmd( + ctx: click.Context, + model: str, + server_url: str | None, + prompt_tokens: int | None, + gen_tokens: int | None, + no_apply: bool, +) -> None: + """Autotune a model's recipe on the actual hardware. + + Sweeps the configurations that matter on Arc — q8_0 KV vs f16 KV with + flash attention forced on (the oneDNN XMX prefill path), across ubatch + sizes — through the running `arc-llama serve`, then persists the fastest. + Requires the server to be running; each candidate costs a model reload. + """ + import asyncio + + from arc_llama.benchmark import ( + TUNE_GEN_TOKENS, + TUNE_PROMPT_TOKENS, + autotune_model, + print_tune_table, + ) + cfg = load_config(ctx.obj["config_path"]) + url = server_url or f"http://{cfg.server.host}:{cfg.server.port}" + try: + outcomes, best = asyncio.run(autotune_model( + url, model, + prompt_tokens=prompt_tokens or TUNE_PROMPT_TOKENS, + gen_tokens=gen_tokens or TUNE_GEN_TOKENS, + apply_best=not no_apply, + cfg=cfg, + )) + except ValueError as e: + console.print(f"[red]{e}[/red]") + sys.exit(1) + except Exception as e: + console.print(f"[red]tune failed: {e}[/red]") + console.print("Is `arc-llama serve` running at " + f"{url}? Start it first — tuning runs through the live server.") + sys.exit(1) + print_tune_table(outcomes, best) + if best is None: + console.print("[yellow]No candidate produced a valid measurement.[/yellow]") + sys.exit(1) + if no_apply: + console.print(f"Winner: [bold]{best.candidate.label}[/bold] (original recipe restored).") + else: + console.print(f"Applied winner: [bold]{best.candidate.label}[/bold] → {best.candidate.edit}") + + # =========================================================================== # mtp-info # =========================================================================== diff --git a/src/arc_llama/config.py b/src/arc_llama/config.py index ecf1ff7..8d53bb1 100644 --- a/src/arc_llama/config.py +++ b/src/arc_llama/config.py @@ -95,6 +95,12 @@ class ServerConfig: admin_token: str | None = None """Bearer token required for destructive admin endpoints (load, stop, edit, scan). Set via `arc-llama serve --admin-token` or env var ARC_LLAMA_ADMIN_TOKEN.""" + jit_cache: str = "managed" + """SYCL JIT cache policy on arches with the persistent-cache SIGSEGV + (Battlemage/Lunar Lake). 'managed' → per-binary fingerprinted cache dir + with crash guard (see sycl_cache.py) — cold starts pay the ~20s JIT cost + once per llama-server build. 'off' → the arch profile's conservative + SYCL_CACHE_PERSISTENT=0, i.e. ~20s JIT on every cold start.""" """Why 11437? Ollama owns 11434 by default, and IPEX-LLM-Ollama installs sometimes use 11435/11436. 11437 is the first port in that neighbourhood that nobody else seems to claim.""" @@ -166,6 +172,10 @@ def launch_recipe(self) -> LaunchRecipe: top_k=r.get("top_k"), spec_type=r.get("spec_type"), ubatch_size=r.get("ubatch_size"), + batch_size=r.get("batch_size"), + flash_attn=r.get("flash_attn"), + cache_reuse=r.get("cache_reuse"), + no_mmap=bool(r.get("no_mmap", False)), extra_flags=list(r.get("extra_flags", [])), ) diff --git a/src/arc_llama/launcher.py b/src/arc_llama/launcher.py index 2025ed8..b9c685b 100644 --- a/src/arc_llama/launcher.py +++ b/src/arc_llama/launcher.py @@ -26,6 +26,7 @@ from arc_llama.arch import Arch, ArchProfile, profile_for from arc_llama.config import Config, GPUConfig, ModelConfig from arc_llama.gguf_meta import has_mtp_heads, is_hybrid_ssm +from arc_llama.sycl_cache import prepare_jit_cache log = logging.getLogger("arc_llama.launcher") @@ -90,9 +91,15 @@ class LaunchPlan: cwd: str | None = None health_url: str = "" backend_url: str = "" + jit_marker: Path | None = None + """Warm-up crash-guard marker for the managed SYCL JIT cache. Written + just before spawn, cleared after the first successful health check or a + clean stop — a leftover marker means the run crashed during JIT warm-up.""" -def build_env(profile: ArchProfile, sycl_index: int) -> dict[str, str]: +def build_env( + profile: ArchProfile, sycl_index: int, extra_env: dict[str, str] | None = None +) -> dict[str, str]: """Compose the environment, layering arch defaults over the user's shell env.""" env = os.environ.copy() # Strip env vars known to break this arch (even if the user inherited them). @@ -102,6 +109,10 @@ def build_env(profile: ArchProfile, sycl_index: int) -> dict[str, str]: # specific GPU index this model is bound to. env.update(profile.sycl_env) env["ONEAPI_DEVICE_SELECTOR"] = f"level_zero:{sycl_index}" + # Layered last so the managed JIT cache can re-enable persistence that the + # arch profile conservatively disabled. + if extra_env: + env.update(extra_env) return env @@ -110,7 +121,29 @@ def build_plan( ) -> LaunchPlan: arch = Arch(gpu.arch) if gpu.arch else Arch.UNKNOWN profile = profile_for(arch) - env = build_env(profile, gpu.sycl_index) + + # --- Managed persistent SYCL JIT cache --- + # Only where the profile would otherwise disable persistence (Xe2), and + # only when the config hasn't opted out. getattr keeps old/partial configs + # (and test doubles) working without a state_dir. + jit_env: dict[str, str] | None = None + jit_marker: Path | None = None + state_dir = getattr(cfg.paths, "state_dir", None) + jit_mode = getattr(cfg.server, "jit_cache", "managed") + if profile.managed_jit_cache and state_dir and jit_mode == "managed": + jit_plan = prepare_jit_cache(state_dir, cfg.paths.llama_server) + if jit_plan.enabled: + jit_env = jit_plan.env + jit_marker = jit_plan.marker + log.info("[%s] %s", model.name, jit_plan.reason) + else: + log.info( + "[%s] managed JIT cache inactive (%s); SYCL_CACHE_PERSISTENT=0 " + "applies — expect ~20s JIT warm-up per cold start", + model.name, jit_plan.reason, + ) + + env = build_env(profile, gpu.sycl_index, extra_env=jit_env) recipe = model.launch_recipe() # --- MTP head detection & safety wiring --- @@ -160,6 +193,7 @@ def build_plan( env=env, backend_url=backend_url, health_url=f"{backend_url}/health", + jit_marker=jit_marker, ) @@ -190,6 +224,14 @@ def start(self, log_dir: Path | None = None) -> None: stdout = self._log_file stderr = subprocess.STDOUT log.info("[%s] starting: %s", self.name, " ".join(self.plan.argv)) + # Arm the JIT crash guard before spawn — a SIGSEGV can land inside the + # first milliseconds of persistent-cache reads. + if self.plan.jit_marker is not None: + try: + self.plan.jit_marker.parent.mkdir(parents=True, exist_ok=True) + self.plan.jit_marker.write_text(f"{os.getpid()} {self.name}\n") + except OSError: + pass self.process = subprocess.Popen( self.plan.argv, env=self.plan.env, @@ -199,6 +241,15 @@ def start(self, log_dir: Path | None = None) -> None: ) self.started_at = time.time() + def _disarm_jit_guard(self) -> None: + """Clear the warm-up marker — the run survived (or was stopped cleanly), + so the persistent JIT cache is not implicated in anything.""" + if self.plan.jit_marker is not None: + try: + self.plan.jit_marker.unlink(missing_ok=True) + except OSError: + pass + async def wait_ready(self, timeout: float = DEFAULT_HEALTH_TIMEOUT) -> bool: deadline = time.time() + timeout async with httpx.AsyncClient(timeout=2.0) as client: @@ -209,6 +260,7 @@ async def wait_ready(self, timeout: float = DEFAULT_HEALTH_TIMEOUT) -> bool: try: r = await client.get(self.plan.health_url) if r.status_code == 200 and r.json().get("status") == "ok": + self._disarm_jit_guard() return True except Exception: pass @@ -216,6 +268,9 @@ async def wait_ready(self, timeout: float = DEFAULT_HEALTH_TIMEOUT) -> bool: return False def stop(self, drain_seconds: float = 3.0) -> None: + # A deliberate stop is not a crash — don't let a Ctrl-C during warm-up + # poison the JIT cache. + self._disarm_jit_guard() if not self.is_running: return proc = self.process diff --git a/src/arc_llama/models.py b/src/arc_llama/models.py index be36b42..f970d98 100644 --- a/src/arc_llama/models.py +++ b/src/arc_llama/models.py @@ -21,7 +21,7 @@ ModelConfig, ) from arc_llama.gguf_meta import has_mtp_heads -from arc_llama.recipes import default_recipe +from arc_llama.recipes import LaunchRecipe, default_recipe log = logging.getLogger("arc_llama.models") @@ -80,6 +80,52 @@ def _short_name_from(repo: str, file: str | None) -> str: return base or "model" +def _recipe_to_dict(recipe: LaunchRecipe) -> dict[str, Any]: + """Serialise the auto-picked recipe fields that belong in config.toml.""" + d: dict[str, Any] = { + "n_gpu_layers": recipe.n_gpu_layers, + "ctx": recipe.ctx, + "parallel": recipe.parallel, + "cache_type_k": recipe.cache_type_k.value, + "cache_type_v": recipe.cache_type_v.value, + } + if recipe.ubatch_size is not None: + d["ubatch_size"] = recipe.ubatch_size + if recipe.batch_size is not None: + d["batch_size"] = recipe.batch_size + if recipe.flash_attn is not None: + d["flash_attn"] = recipe.flash_attn + if recipe.cache_reuse is not None: + d["cache_reuse"] = recipe.cache_reuse + return d + + +def warn_slow_weight_quant(arch, filename: str) -> str | None: + """Warn when a GGUF's weight quant hits a known-slow kernel path. + + Q8_0 on Xe2 reaches ~22% of memory bandwidth vs Q4_K_M's ~55% + (ggml-org/llama.cpp#21517) — 4-5x slower generation for 1.7x the bytes. + Returns the warning string (also logged) or None. + """ + from arc_llama.arch import profile_for + profile = profile_for(arch) + if not profile.slow_weight_quants: + return None + m = _QUANT_TIER_RE.search(filename) + if m is None: + return None + tier = m.group(1).upper() + if tier not in profile.slow_weight_quants: + return None + msg = ( + f"{filename}: {tier} weight quant has a known-slow kernel path on " + f"{profile.display_name} — expect 4-5x slower generation than Q4_K_M. " + f"Prefer Q4_K_M, Q4_K_S, or Q6_K of the same model." + ) + log.warning(msg) + return msg + + def add_local_model( cfg: Config, *, @@ -121,17 +167,16 @@ def add_local_model( model_file_mb=p.stat().st_size // (1024 * 1024), kv_class=kv_class, ) - recipe_dict: dict[str, Any] = { - "n_gpu_layers": recipe.n_gpu_layers, - "ctx": recipe.ctx, - "parallel": recipe.parallel, - "cache_type_k": recipe.cache_type_k.value, - "cache_type_v": recipe.cache_type_v.value, - } + warn_slow_weight_quant(arch, p.name) + recipe_dict = _recipe_to_dict(recipe) # Auto-enable draft-mtp for models that actually carry MTP heads. + # ubatch is pinned to 8 there, so the XMX prefill batch/FA fields would be + # dead weight (and -fa on is untested against hybrid-SSM MTP graphs). if has_mtp_heads(p): recipe_dict["spec_type"] = "draft-mtp" recipe_dict["ubatch_size"] = 8 + recipe_dict.pop("batch_size", None) + recipe_dict.pop("flash_attn", None) log.info( "model %s has MTP heads; auto-enabling spec_type=draft-mtp, ubatch_size=8", name, @@ -313,17 +358,15 @@ def register_discovered( model_file_mb=rp.stat().st_size // (1024 * 1024), kv_class=kv_class, ) - recipe_dict: dict[str, Any] = { - "n_gpu_layers": recipe.n_gpu_layers, - "ctx": recipe.ctx, - "parallel": recipe.parallel, - "cache_type_k": recipe.cache_type_k.value, - "cache_type_v": recipe.cache_type_v.value, - } + warn_slow_weight_quant(arch, rp.name) + recipe_dict = _recipe_to_dict(recipe) # Auto-enable draft-mtp for discovered models that carry MTP heads. + # Same field trimming as add_local_model — ub is pinned to 8. if has_mtp_heads(rp): recipe_dict["spec_type"] = "draft-mtp" recipe_dict["ubatch_size"] = 8 + recipe_dict.pop("batch_size", None) + recipe_dict.pop("flash_attn", None) log.info( "discovered %s has MTP heads; auto-enabling spec_type=draft-mtp, ubatch_size=8", rp.name, diff --git a/src/arc_llama/recipes.py b/src/arc_llama/recipes.py index 306d4f6..245a7a9 100644 --- a/src/arc_llama/recipes.py +++ b/src/arc_llama/recipes.py @@ -51,6 +51,19 @@ class LaunchRecipe: """Speculative decoding type, e.g. 'draft-mtp'.""" ubatch_size: int | None = None """Ubatch size (-ub). Auto-set to 8 for MTP models to avoid SSM compute-buffer OOM.""" + batch_size: int | None = None + """Logical batch size (-b). Raised together with ubatch for prefill throughput.""" + flash_attn: str | None = None + """Flash attention (-fa): 'on', 'off', or 'auto'. None → omit the flag and + let llama.cpp decide. Must be 'on' + f16 KV cache to hit the oneDNN XMX + prefill path on Xe2.""" + cache_reuse: int | None = None + """--cache-reuse N: reuse cached prompt prefixes with a KV shift when at + least N tokens match. Slashes time-to-first-token in multi-turn chat. + Incompatible with sliding-window-attention models (gemma_swa).""" + no_mmap: bool = False + """--no-mmap: read the model up front instead of paging it in lazily. + Avoids first-inference page-fault stalls when weights live on slow disks.""" extra_flags: list[str] = field(default_factory=list) """Anything else the user wants appended to the command line verbatim.""" @@ -74,6 +87,14 @@ def to_argv(self) -> list[str]: argv += ["--spec-type", self.spec_type] if self.ubatch_size is not None: argv += ["-ub", str(self.ubatch_size)] + if self.batch_size is not None: + argv += ["-b", str(self.batch_size)] + if self.flash_attn is not None: + argv += ["-fa", self.flash_attn] + if self.cache_reuse is not None: + argv += ["--cache-reuse", str(self.cache_reuse)] + if self.no_mmap: + argv += ["--no-mmap"] argv += list(self.extra_flags) return argv @@ -134,6 +155,25 @@ def suggest_ctx( return max(4096, min(rounded, ctx_cap)) +XMX_FA_MIN_F16_CTX = 16384 +"""Smallest f16-KV context worth trading Q8_0 KV headroom for. Below this the +KV savings buy more useful context than the XMX prefill path buys speed.""" + +XMX_FA_UBATCH = 1024 +XMX_FA_BATCH = 2048 +"""Prefill batch sizing for the XMX path. llama.cpp's default -ub 512 leaves +the systolic arrays underfed on Arc dGPUs; 1024/2048 measured fastest on +Battlemage without blowing up compute buffers.""" + +XMX_FA_COMPUTE_BUFFER_MB = 1536 +"""Compute-buffer estimate when ubatch is raised to 1024 — roughly double the +768 MB default-ubatch figure. Deliberately pessimistic so ctx sizing stays safe.""" + +DEFAULT_CACHE_REUSE = 256 +"""--cache-reuse threshold applied to non-SWA models: multi-turn chats re-use +the common prefix KV instead of re-prefilling the whole conversation.""" + + def default_recipe( arch: Arch, vram_mb: int, @@ -141,8 +181,38 @@ def default_recipe( kv_class: str = "default", prefer_q8_kv: bool = True, ) -> LaunchRecipe: - """A safe starting recipe for a freshly added model on a given arch.""" + """A safe starting recipe for a freshly added model on a given arch. + + On Xe2 (Battlemage / Lunar Lake) the oneDNN XMX flash-attention path is + preferred when VRAM allows: it needs an f16 KV cache (quantized KV falls + back to shader kernels), so we take f16 + `-fa on` whenever that still + leaves ≥ XMX_FA_MIN_F16_CTX of context, and only drop to q8_0 KV on + VRAM-tight setups where context matters more than prefill speed. + """ profile: ArchProfile = profile_for(arch) + cache_reuse = None if kv_class == "gemma_swa" else DEFAULT_CACHE_REUSE + + if profile.supports_xmx_fa: + ctx_f16 = suggest_ctx( + vram_mb=vram_mb, + model_file_mb=model_file_mb, + kv_type=KVCacheType.F16, + kv_class=kv_class, + compute_buffer_mb=XMX_FA_COMPUTE_BUFFER_MB, + ) + if ctx_f16 >= XMX_FA_MIN_F16_CTX: + return LaunchRecipe( + n_gpu_layers=999, + ctx=ctx_f16, + parallel=1, + cache_type_k=KVCacheType.F16, + cache_type_v=KVCacheType.F16, + ubatch_size=XMX_FA_UBATCH, + batch_size=XMX_FA_BATCH, + flash_attn="on", + cache_reuse=cache_reuse, + ) + kv_type = KVCacheType.Q8_0 if (prefer_q8_kv and profile.safe_kv_q8) else KVCacheType.F16 ctx = suggest_ctx( vram_mb=vram_mb, @@ -156,4 +226,5 @@ def default_recipe( parallel=1, cache_type_k=kv_type, cache_type_v=kv_type, + cache_reuse=cache_reuse, ) diff --git a/src/arc_llama/server.py b/src/arc_llama/server.py index f902385..8c44364 100644 --- a/src/arc_llama/server.py +++ b/src/arc_llama/server.py @@ -476,7 +476,9 @@ async def admin_edit_model(name: str, request: Request) -> dict: Body is a partial recipe dict — only provided fields change. Recognised fields: `ctx`, `cache_type_k`, `cache_type_v`, `parallel`, `kv_class`, - `spec_type`, `ubatch_size`. + `spec_type`, `ubatch_size`, `batch_size`, `flash_attn`, `cache_reuse`. + `batch_size`, `flash_attn`, and `cache_reuse` accept null to clear the + field (falls back to llama.cpp defaults). If the model is currently loaded, the server is stopped first; callers decide whether to reload it afterwards via /admin/load. """ @@ -542,14 +544,52 @@ async def admin_edit_model(name: str, request: Request) -> dict: recipe["spec_type"] = v changed.append("spec_type") if "ubatch_size" in body: - try: - ub = int(body["ubatch_size"]) - except (TypeError, ValueError): - raise HTTPException(status_code=400, detail="ubatch_size must be an integer") from None - if not (1 <= ub <= 4096): - raise HTTPException(status_code=400, detail="ubatch_size must be 1..4096") - recipe["ubatch_size"] = ub + if body["ubatch_size"] is None: + recipe.pop("ubatch_size", None) + else: + try: + ub = int(body["ubatch_size"]) + except (TypeError, ValueError): + raise HTTPException(status_code=400, detail="ubatch_size must be an integer") from None + if not (1 <= ub <= 4096): + raise HTTPException(status_code=400, detail="ubatch_size must be 1..4096") + recipe["ubatch_size"] = ub changed.append("ubatch_size") + if "batch_size" in body: + if body["batch_size"] is None: + recipe.pop("batch_size", None) + else: + try: + b = int(body["batch_size"]) + except (TypeError, ValueError): + raise HTTPException(status_code=400, detail="batch_size must be an integer") from None + if not (1 <= b <= 8192): + raise HTTPException(status_code=400, detail="batch_size must be 1..8192") + recipe["batch_size"] = b + changed.append("batch_size") + if "flash_attn" in body: + if body["flash_attn"] is None: + recipe.pop("flash_attn", None) + else: + v = str(body["flash_attn"]) + if v not in ("on", "off", "auto"): + raise HTTPException( + status_code=400, detail="flash_attn must be 'on', 'off', or 'auto'", + ) + recipe["flash_attn"] = v + changed.append("flash_attn") + if "cache_reuse" in body: + if body["cache_reuse"] is None: + recipe.pop("cache_reuse", None) + else: + try: + cr = int(body["cache_reuse"]) + except (TypeError, ValueError): + raise HTTPException(status_code=400, detail="cache_reuse must be an integer") from None + if not (0 <= cr <= 65536): + raise HTTPException(status_code=400, detail="cache_reuse must be 0..65536") + recipe["cache_reuse"] = cr + changed.append("cache_reuse") if not changed: raise HTTPException(status_code=400, detail="no recognised fields to edit") model.recipe = recipe diff --git a/src/arc_llama/sycl_cache.py b/src/arc_llama/sycl_cache.py new file mode 100644 index 0000000..a4ae0f7 --- /dev/null +++ b/src/arc_llama/sycl_cache.py @@ -0,0 +1,181 @@ +"""Managed persistent SYCL JIT cache. + +Battlemage with libsycl.so.9 from oneAPI 2026.0 reproducibly SIGSEGVs in +`PersistentDeviceCodeCache::getItemFromDisc` when `SYCL_CACHE_PERSISTENT=1` +reads back cache entries written by a *different* stack (older llama-server +build, different driver, interrupted write). The stock workaround — disabling +the cache outright — costs ~20 s of JIT recompilation on every cold start. + +This module keeps persistence instead of giving it up, made safe two ways: + +1. **Fingerprint isolation.** Every (llama-server binary, kernel GPU driver) + combination gets its own private cache directory under + `/sycl-cache/`. A cache is only ever read by the + exact stack that wrote it, so the stale-entry crash has nothing to bite on. + Upgrading llama.cpp or the driver simply lands in a fresh directory; old + ones are pruned. + +2. **Crash guard.** The launcher drops a `warming` marker in the cache dir + before spawning llama-server and removes it once the health check passes + (or on a clean stop). If we ever find a leftover marker, the previous run + died mid-warm-up — we wipe the directory, poison that fingerprint, and fall + back to `SYCL_CACHE_PERSISTENT=0`. Worst case is exactly the behaviour + arc-llama shipped with; best case (the common one) cold starts pay the JIT + cost once per llama-server build instead of every time. +""" +from __future__ import annotations + +import hashlib +import logging +import shutil +import time +from dataclasses import dataclass, field +from pathlib import Path + +log = logging.getLogger("arc_llama.sycl_cache") + +POISON_FILE = "POISONED" +MARKER_FILE = "warming" +KEEP_RECENT_CACHES = 2 +"""Old fingerprint dirs kept alongside the active one (rollbacks are cheap).""" + +_DRIVER_VERSION_PROBES = ( + Path("/sys/module/xe/srcversion"), + Path("/sys/module/i915/srcversion"), +) + + +@dataclass +class JitCachePlan: + """Outcome of preparing the managed cache for one launch.""" + enabled: bool + reason: str + env: dict[str, str] = field(default_factory=dict) + """Env overrides to layer on top of the arch profile (may re-enable + SYCL_CACHE_PERSISTENT that the profile disabled).""" + marker: Path | None = None + """Warm-up marker the launcher must write before spawn and clear after + the first successful health check.""" + cache_dir: Path | None = None + + +def binary_fingerprint(llama_server: str) -> str | None: + """Fingerprint the llama-server binary + GPU driver combination. + + Uses (resolved path, size, mtime) rather than hashing the multi-hundred-MB + binary; a rebuild or upgrade always changes at least one of those. Returns + None if the binary can't be found — callers must fall back to disabled. + """ + exe = shutil.which(llama_server) or llama_server + p = Path(exe).expanduser() + try: + p = p.resolve() + st = p.stat() + except OSError: + return None + driver = "" + for probe in _DRIVER_VERSION_PROBES: + try: + driver += probe.read_text().strip() + except OSError: + continue + raw = f"{p}|{st.st_size}|{st.st_mtime_ns}|{driver}" + return hashlib.sha256(raw.encode()).hexdigest()[:16] + + +def _wipe_cache_entries(cache_dir: Path) -> None: + """Remove everything in the dir except the poison file.""" + try: + for entry in cache_dir.iterdir(): + if entry.name == POISON_FILE: + continue + if entry.is_dir(): + shutil.rmtree(entry, ignore_errors=True) + else: + entry.unlink(missing_ok=True) + except OSError: + pass + + +def _prune_old_caches(root: Path, current: str) -> None: + """Delete fingerprint dirs beyond the KEEP_RECENT_CACHES most recent.""" + try: + siblings = [ + d for d in root.iterdir() + if d.is_dir() and d.name != current + ] + except OSError: + return + siblings.sort(key=lambda d: d.stat().st_mtime, reverse=True) + for stale in siblings[KEEP_RECENT_CACHES:]: + log.info("pruning stale SYCL JIT cache %s", stale) + shutil.rmtree(stale, ignore_errors=True) + + +def prepare_jit_cache(state_dir: str | Path, llama_server: str) -> JitCachePlan: + """Prepare the managed cache dir for a launch and return env + marker. + + Call once per llama-server spawn. Never raises — any filesystem trouble + degrades to a disabled plan, which leaves the arch profile's conservative + `SYCL_CACHE_PERSISTENT=0` in effect. + """ + fp = binary_fingerprint(llama_server) + if fp is None: + return JitCachePlan( + enabled=False, + reason=f"llama-server binary not found for fingerprinting: {llama_server}", + ) + root = Path(state_dir).expanduser() / "sycl-cache" + cache_dir = root / fp + poison = cache_dir / POISON_FILE + marker = cache_dir / MARKER_FILE + + if poison.exists(): + return JitCachePlan( + enabled=False, + reason="fingerprint poisoned by an earlier warm-up crash " + "(delete the POISONED file to retry)", + cache_dir=cache_dir, + ) + + if marker.exists(): + # Previous run died between spawn and first health check with this + # cache active. Assume the persistent cache is implicated: wipe it and + # never re-enable for this fingerprint. A user Ctrl-C during warm-up + # also lands here — deliberately conservative, and self-documenting on + # disk via the poison file. + log.warning( + "leftover warm-up marker in %s — previous run crashed during SYCL " + "JIT warm-up; wiping and poisoning this cache", cache_dir, + ) + _wipe_cache_entries(cache_dir) + try: + poison.write_text( + f"poisoned {time.strftime('%Y-%m-%dT%H:%M:%S%z')}: previous " + "llama-server run crashed before its first health check while " + "this persistent JIT cache was enabled.\n" + ) + except OSError: + pass + return JitCachePlan( + enabled=False, + reason="previous run crashed during JIT warm-up; cache wiped and poisoned", + cache_dir=cache_dir, + ) + + try: + cache_dir.mkdir(parents=True, exist_ok=True) + except OSError as e: + return JitCachePlan(enabled=False, reason=f"cannot create cache dir: {e}") + _prune_old_caches(root, current=fp) + + return JitCachePlan( + enabled=True, + reason=f"managed persistent JIT cache at {cache_dir}", + env={ + "SYCL_CACHE_PERSISTENT": "1", + "SYCL_CACHE_DIR": str(cache_dir), + }, + marker=marker, + cache_dir=cache_dir, + ) diff --git a/tests/test_launcher.py b/tests/test_launcher.py index 24a8544..0927893 100644 --- a/tests/test_launcher.py +++ b/tests/test_launcher.py @@ -225,3 +225,51 @@ def test_stop_idempotent(self): # Should not raise when not running srv.stop() assert srv.is_running is False + + +class TestManagedJitCache: + def _cfg(self, tmp_path: Path) -> Config: + binary = tmp_path / "llama-server" + binary.write_bytes(b"ELF fake") + cfg = Config() + cfg.paths.llama_server = str(binary) + cfg.paths.state_dir = str(tmp_path / "state") + return cfg + + def test_battlemage_gets_managed_cache(self, tmp_path: Path): + cfg = self._cfg(tmp_path) + model = ModelConfig(name="m", path="/m.gguf", port=18080, gpu_pci_slot="00:00.0") + gpu = GPUConfig(pci_slot="00:00.0", sycl_index=0, arch="battlemage") + plan = build_plan(cfg, model, gpu) + # Managed cache overrides the profile's conservative disable. + assert plan.env["SYCL_CACHE_PERSISTENT"] == "1" + assert "sycl-cache" in plan.env["SYCL_CACHE_DIR"] + assert plan.jit_marker is not None + + def test_jit_cache_off_keeps_profile_default(self, tmp_path: Path): + cfg = self._cfg(tmp_path) + cfg.server.jit_cache = "off" + model = ModelConfig(name="m", path="/m.gguf", port=18080, gpu_pci_slot="00:00.0") + gpu = GPUConfig(pci_slot="00:00.0", sycl_index=0, arch="battlemage") + plan = build_plan(cfg, model, gpu) + assert plan.env["SYCL_CACHE_PERSISTENT"] == "0" + assert plan.jit_marker is None + + def test_alchemist_untouched(self, tmp_path: Path): + cfg = self._cfg(tmp_path) + model = ModelConfig(name="m", path="/m.gguf", port=18080, gpu_pci_slot="00:00.0") + gpu = GPUConfig(pci_slot="00:00.0", sycl_index=0, arch="alchemist") + plan = build_plan(cfg, model, gpu) + assert "SYCL_CACHE_DIR" not in plan.env + assert plan.jit_marker is None + + def test_stop_disarms_guard(self, tmp_path: Path): + cfg = self._cfg(tmp_path) + model = ModelConfig(name="m", path="/m.gguf", port=18080, gpu_pci_slot="00:00.0") + gpu = GPUConfig(pci_slot="00:00.0", sycl_index=0, arch="battlemage") + plan = build_plan(cfg, model, gpu) + plan.jit_marker.parent.mkdir(parents=True, exist_ok=True) + plan.jit_marker.write_text("pid\n") + srv = LlamaServer(plan) + srv.stop() # clean stop, not running — must still clear the marker + assert not plan.jit_marker.exists() diff --git a/tests/test_models.py b/tests/test_models.py index a290fa6..d4e2059 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -680,3 +680,44 @@ def test_discover_and_register_ggufs_skips_hidden_symlink_and_existing_files(tmp assert added[0].kv_class == "qwen3_27b_dense" assert added_again == [] assert len(cfg.models) == 1 + + +# =========================================================================== +# slow weight-quant warning + new recipe field persistence +# =========================================================================== + + +class TestSlowWeightQuantWarning: + def test_q8_0_on_battlemage_warns(self): + from arc_llama.arch import Arch + from arc_llama.models import warn_slow_weight_quant + msg = warn_slow_weight_quant(Arch.BATTLEMAGE, "Qwen3.6-27B-Q8_0.gguf") + assert msg is not None + assert "Q8_0" in msg + + def test_q4_k_m_on_battlemage_silent(self): + from arc_llama.arch import Arch + from arc_llama.models import warn_slow_weight_quant + assert warn_slow_weight_quant(Arch.BATTLEMAGE, "Qwen3.6-27B-Q4_K_M.gguf") is None + + def test_q8_0_on_alchemist_silent(self): + # No measured slow path on Alchemist — don't cry wolf. + from arc_llama.arch import Arch + from arc_llama.models import warn_slow_weight_quant + assert warn_slow_weight_quant(Arch.ALCHEMIST, "model-Q8_0.gguf") is None + + +def test_add_local_model_persists_xmx_fields(tmp_path): + """Real default_recipe on a roomy Battlemage → FA fields land in config.""" + with patch("arc_llama.models.has_mtp_heads", return_value=False): + cfg = _make_config_with_gpu(tmp_path) + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"\x00" * 1024) + mc = add_local_model( + cfg, name="m", path=str(gguf), gpu_pci_slot="0000:03:00.0", + ) + assert mc.recipe["cache_type_k"] == "f16" + assert mc.recipe["flash_attn"] == "on" + assert mc.recipe["ubatch_size"] == 1024 + assert mc.recipe["batch_size"] == 2048 + assert mc.recipe["cache_reuse"] == 256 diff --git a/tests/test_recipes.py b/tests/test_recipes.py index c05a0af..f209c22 100644 --- a/tests/test_recipes.py +++ b/tests/test_recipes.py @@ -96,15 +96,27 @@ def test_q4_saves_more_than_q8(self): class TestDefaultRecipe: - def test_battlemage_prefers_q8(self): + def test_battlemage_roomy_vram_prefers_f16_xmx_fa(self): + # Since the oneDNN XMX FA path (f16 KV only), roomy Battlemage + # setups take f16 KV + forced FA instead of the old q8_0 default. r = default_recipe( Arch.BATTLEMAGE, vram_mb=24 * 1024, model_file_mb=4 * 1024, ) + assert r.cache_type_k == KVCacheType.F16 + assert r.cache_type_v == KVCacheType.F16 + assert r.flash_attn == "on" + assert r.n_gpu_layers == 999 + + def test_battlemage_tight_vram_still_prefers_q8(self): + r = default_recipe( + Arch.BATTLEMAGE, + vram_mb=8 * 1024, + model_file_mb=6 * 1024, + ) assert r.cache_type_k == KVCacheType.Q8_0 assert r.cache_type_v == KVCacheType.Q8_0 - assert r.n_gpu_layers == 999 def test_unknown_arch_is_conservative(self): r = default_recipe( @@ -192,3 +204,71 @@ def test_spec_type_and_ubatch_size(self): assert argv[argv.index("--spec-type") + 1] == "draft-mtp" assert "-ub" in argv assert argv[argv.index("-ub") + 1] == "8" + + +class TestNewLaunchFlags: + def test_xmx_flags_emitted(self): + r = LaunchRecipe( + flash_attn="on", + ubatch_size=1024, + batch_size=2048, + cache_reuse=256, + no_mmap=True, + ) + argv = r.to_argv() + assert argv[argv.index("-fa") + 1] == "on" + assert argv[argv.index("-ub") + 1] == "1024" + assert argv[argv.index("-b") + 1] == "2048" + assert argv[argv.index("--cache-reuse") + 1] == "256" + assert "--no-mmap" in argv + + def test_new_flags_omitted_by_default(self): + argv = LaunchRecipe().to_argv() + assert "-fa" not in argv + assert "-b" not in argv + assert "--cache-reuse" not in argv + assert "--no-mmap" not in argv + + +class TestXmxFaPolicy: + def test_battlemage_roomy_vram_prefers_f16_fa(self): + """24GB card, 4GB model → f16 KV + forced FA + fat prefill batches.""" + r = default_recipe( + Arch.BATTLEMAGE, vram_mb=24 * 1024, model_file_mb=4 * 1024, + ) + assert r.cache_type_k == KVCacheType.F16 + assert r.cache_type_v == KVCacheType.F16 + assert r.flash_attn == "on" + assert r.ubatch_size == 1024 + assert r.batch_size == 2048 + assert r.cache_reuse == 256 + + def test_battlemage_tight_vram_falls_back_to_q8(self): + """8GB card, 6.5GB model → f16 ctx would be tiny; keep q8_0 KV.""" + r = default_recipe( + Arch.BATTLEMAGE, vram_mb=8 * 1024, model_file_mb=6656, + ) + assert r.cache_type_k == KVCacheType.Q8_0 + assert r.flash_attn is None + assert r.ubatch_size is None + + def test_alchemist_unchanged_q8_default(self): + """No confirmed XMX FA path on Alchemist — behaviour is pre-existing.""" + r = default_recipe( + Arch.ALCHEMIST, vram_mb=16 * 1024, model_file_mb=4 * 1024, + ) + assert r.cache_type_k == KVCacheType.Q8_0 + assert r.flash_attn is None + + def test_gemma_swa_gets_no_cache_reuse(self): + r = default_recipe( + Arch.BATTLEMAGE, vram_mb=24 * 1024, model_file_mb=4 * 1024, + kv_class="gemma_swa", + ) + assert r.cache_reuse is None + + def test_non_swa_gets_cache_reuse_on_fallback_path(self): + r = default_recipe( + Arch.ALCHEMIST, vram_mb=16 * 1024, model_file_mb=4 * 1024, + ) + assert r.cache_reuse == 256 diff --git a/tests/test_sycl_cache.py b/tests/test_sycl_cache.py new file mode 100644 index 0000000..715759b --- /dev/null +++ b/tests/test_sycl_cache.py @@ -0,0 +1,107 @@ +"""Tests for arc_llama.sycl_cache — fingerprinting, crash guard, pruning.""" +from __future__ import annotations + +import time +from pathlib import Path + +from arc_llama.sycl_cache import ( + KEEP_RECENT_CACHES, + MARKER_FILE, + POISON_FILE, + binary_fingerprint, + prepare_jit_cache, +) + + +def _fake_binary(tmp_path: Path, name: str = "llama-server", content: bytes = b"ELF") -> Path: + p = tmp_path / name + p.write_bytes(content) + return p + + +class TestFingerprint: + def test_missing_binary_returns_none(self, tmp_path: Path): + assert binary_fingerprint(str(tmp_path / "nope")) is None + + def test_stable_for_same_file(self, tmp_path: Path): + b = _fake_binary(tmp_path) + assert binary_fingerprint(str(b)) == binary_fingerprint(str(b)) + + def test_changes_when_binary_changes(self, tmp_path: Path): + b = _fake_binary(tmp_path) + fp1 = binary_fingerprint(str(b)) + # New size + mtime — simulates a llama.cpp upgrade in place. + b.write_bytes(b"ELF v2 rebuilt") + fp2 = binary_fingerprint(str(b)) + assert fp1 != fp2 + + +class TestPrepare: + def test_enables_and_creates_dir(self, tmp_path: Path): + b = _fake_binary(tmp_path) + plan = prepare_jit_cache(tmp_path / "state", str(b)) + assert plan.enabled is True + assert plan.env["SYCL_CACHE_PERSISTENT"] == "1" + assert plan.cache_dir is not None and plan.cache_dir.is_dir() + assert plan.env["SYCL_CACHE_DIR"] == str(plan.cache_dir) + assert plan.marker == plan.cache_dir / MARKER_FILE + + def test_missing_binary_disables(self, tmp_path: Path): + plan = prepare_jit_cache(tmp_path / "state", str(tmp_path / "nope")) + assert plan.enabled is False + assert plan.env == {} + + def test_leftover_marker_wipes_and_poisons(self, tmp_path: Path): + b = _fake_binary(tmp_path) + first = prepare_jit_cache(tmp_path / "state", str(b)) + assert first.enabled + # Simulate a crash during warm-up: marker written, never cleared, + # plus a cache entry the runtime wrote before dying. + first.marker.write_text("12345 test\n") + (first.cache_dir / "some-kernel.bin").write_bytes(b"jit blob") + + second = prepare_jit_cache(tmp_path / "state", str(b)) + assert second.enabled is False + assert (first.cache_dir / POISON_FILE).exists() + assert not (first.cache_dir / "some-kernel.bin").exists() + assert not first.marker.exists() + + def test_poison_is_sticky(self, tmp_path: Path): + b = _fake_binary(tmp_path) + first = prepare_jit_cache(tmp_path / "state", str(b)) + first.marker.write_text("crash\n") + prepare_jit_cache(tmp_path / "state", str(b)) # poisons + third = prepare_jit_cache(tmp_path / "state", str(b)) + assert third.enabled is False + assert "poisoned" in third.reason + + def test_clean_lifecycle_not_poisoned(self, tmp_path: Path): + """Marker written then cleared (healthy run) → next prepare re-enables.""" + b = _fake_binary(tmp_path) + first = prepare_jit_cache(tmp_path / "state", str(b)) + first.marker.write_text("pid\n") + first.marker.unlink() + second = prepare_jit_cache(tmp_path / "state", str(b)) + assert second.enabled is True + assert second.cache_dir == first.cache_dir + + def test_prunes_old_fingerprints(self, tmp_path: Path): + state = tmp_path / "state" + root = state / "sycl-cache" + root.mkdir(parents=True) + # Simulate caches left behind by old llama-server builds. + stale = [] + for i in range(KEEP_RECENT_CACHES + 3): + d = root / f"oldfp{i:02d}" + d.mkdir() + t = time.time() - 1000 + i + import os + os.utime(d, (t, t)) + stale.append(d) + b = _fake_binary(tmp_path) + plan = prepare_jit_cache(state, str(b)) + assert plan.enabled + survivors = [d for d in stale if d.exists()] + assert len(survivors) == KEEP_RECENT_CACHES + # The newest siblings survive. + assert survivors == stale[-KEEP_RECENT_CACHES:]