From 9cd91725a365fe2bb9fa034ae7441f6610d98bce Mon Sep 17 00:00:00 2001 From: szibis <329831+szibis@users.noreply.github.com> Date: Mon, 11 May 2026 14:13:00 +0200 Subject: [PATCH 1/6] =?UTF-8?q?feat:=20v0.9.0=20=E2=80=94=20SSD=20KV=20cac?= =?UTF-8?q?he,=20prefix=20trie,=20telemetry,=20DFlash=20improvements?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SSD KV cache persistence: 2-tier hot/cold with safetensors serialization to disk. Evicts oldest entries to SSD when RAM threshold exceeded, restores on cache miss. Session-isolated via prompt prefix hash. Prefix caching: RadixAttention-style compressed trie for KV cache sharing across sessions. Longest-prefix matching with LRU eviction. Thread-safe. Hardware telemetry: IOReport-based sampling for GPU utilization, power draw, CPU/GPU temperature, ANE utilization. Background daemon thread with circular buffer history. New /telemetry endpoints in Python server and Rust proxy. Dashboard: power consumption chart, temperature chart (CPU/GPU with color coding), ANE utilization display. Polls /telemetry/current alongside existing status endpoints. DFlash acceptance improvements: soft threshold acceptance, hidden state normalization to bridge quantized/fp16 gap, proper top-K verification in cached path, block size auto-tuning based on measured acceptance rate. Expert prefetch from drafts: predicts MoE expert assignments from draft tokens via router gate weights, triggers async cache prefetch before target verification. DeepSeek V4 Flash: added mini preset (3-layer, 1024-hidden), auto preset selector based on available RAM, custom calibration prompts via --calibration-prompts. 1995 tests passing, 0 failures. --- mlx-flash-server/src/dashboard.rs | 78 ++++- mlx-flash-server/src/server.rs | 40 +++ mlx_flash_compress/__init__.py | 16 + mlx_flash_compress/dflash_model.py | 225 +++++++++--- mlx_flash_compress/draft_expert_prefetch.py | 150 ++++++++ mlx_flash_compress/kv_cache_backend.py | 20 +- mlx_flash_compress/prefix_cache.py | 268 +++++++++++++++ mlx_flash_compress/serve.py | 38 +++ mlx_flash_compress/ssd_kv_cache.py | 199 +++++++++++ mlx_flash_compress/telemetry.py | 340 ++++++++++++++++++ scripts/train_dflash_drafter.py | 209 +++++++++--- tests/test_dflash_acceptance.py | 359 ++++++++++++++++++++ tests/test_draft_expert_prefetch.py | 201 +++++++++++ tests/test_imports.py | 7 + tests/test_prefix_cache.py | 325 ++++++++++++++++++ tests/test_ssd_kv_cache.py | 313 +++++++++++++++++ tests/test_telemetry.py | 330 ++++++++++++++++++ 17 files changed, 3023 insertions(+), 95 deletions(-) create mode 100644 mlx_flash_compress/draft_expert_prefetch.py create mode 100644 mlx_flash_compress/prefix_cache.py create mode 100644 mlx_flash_compress/ssd_kv_cache.py create mode 100644 mlx_flash_compress/telemetry.py create mode 100644 tests/test_dflash_acceptance.py create mode 100644 tests/test_draft_expert_prefetch.py create mode 100644 tests/test_prefix_cache.py create mode 100644 tests/test_ssd_kv_cache.py create mode 100644 tests/test_telemetry.py diff --git a/mlx-flash-server/src/dashboard.rs b/mlx-flash-server/src/dashboard.rs index 7056935..97a8a0b 100644 --- a/mlx-flash-server/src/dashboard.rs +++ b/mlx-flash-server/src/dashboard.rs @@ -184,6 +184,7 @@ const DASHBOARD_HTML: &str = r##"
-- GB used
--% renderer / --% tiler
+
Power: -- W | ANE: --%
Optimization Hints
@@ -200,7 +201,23 @@ const DASHBOARD_HTML: &str = r##"
- + +
+
+
Power (W)
+
0 W
+ +
+
+
+
+
Temperature (°C)
+
-- °C
+ +
+
+ +
Live Logs last 100 entries
@@ -212,16 +229,22 @@ const DASHBOARD_HTML: &str = r##" """ + def _handle_telemetry(self): + """Return current telemetry sample plus 120-sample history.""" + state = self.server_state + stats = state.telemetry.get_stats() + history = state.telemetry.get_history(seconds=120) + self._send_json( + { + "current": stats.get("current", {}), + "stats": { + "samples_count": stats.get("samples_count", 0), + "avg": stats.get("avg", {}), + "max": stats.get("max", {}), + "min": stats.get("min", {}), + }, + "history": history, + } + ) + + def _handle_telemetry_current(self): + """Return just the current telemetry sample.""" + state = self.server_state + sample = state.telemetry.sample() + self._send_json(sample.to_dict()) + def _send_json(self, data: dict, status: int = 200): self.send_response(status) self.send_header("Content-Type", "application/json") @@ -1327,6 +1359,10 @@ def main(): ChatHandler.server_state = state + # Start hardware telemetry sampling + state.telemetry.start_sampling(interval_ms=1000) + logger.info("Hardware telemetry started", extra={"action": "telemetry_start", "interval_ms": 1000}) + server = ThreadedHTTPServer((args.host, args.port), ChatHandler) logger.info( "Server listening", @@ -1342,6 +1378,7 @@ def _graceful_shutdown(signum, frame): logger.info( f"Received {sig_name}, shutting down gracefully", extra={"action": "server_stop", "signal": sig_name} ) + state.telemetry.stop_sampling() server.shutdown() signal.signal(signal.SIGTERM, _graceful_shutdown) @@ -1351,6 +1388,7 @@ def _graceful_shutdown(signum, frame): server.serve_forever() except KeyboardInterrupt: logger.info("Shutting down", extra={"action": "server_stop"}) + state.telemetry.stop_sampling() server.shutdown() diff --git a/mlx_flash_compress/ssd_kv_cache.py b/mlx_flash_compress/ssd_kv_cache.py new file mode 100644 index 0000000..0add479 --- /dev/null +++ b/mlx_flash_compress/ssd_kv_cache.py @@ -0,0 +1,199 @@ +"""2-tier KV cache: hot in RAM, cold on SSD. + +When the RAM cache exceeds ``max_ram_tokens`` the oldest entries are +evicted to SSD as safetensors files. On a cache miss the cold entries +are loaded back and merged with the hot portion, avoiding recomputation. + +Each inference session gets its own SSD subdirectory keyed by a hash of +the prompt prefix, so concurrent sessions never collide. +""" + +from __future__ import annotations + +import hashlib +import os +import shutil +import threading +import time +from pathlib import Path +from typing import Any + +import mlx.core as mx + +from mlx_flash_compress.kv_cache_backend import KVCacheBackend, PlainKVCache + + +def _session_dir(base: Path, prompt_prefix: str) -> Path: + h = hashlib.sha256(prompt_prefix.encode()).hexdigest()[:16] + return base / h + + +class SSDKVCache(KVCacheBackend): + """Hot/cold KV cache with SSD spill. + + Parameters + ---------- + num_layers: + Number of transformer layers. + num_heads: + Number of KV attention heads. + head_dim: + Dimension per head. + max_ram_tokens: + Maximum tokens kept in the hot (RAM) tier per layer before + eviction to SSD. + cache_dir: + Root directory for cold storage. A session subdirectory is + created underneath. + prompt_prefix: + Used to derive a session-unique subdirectory hash. + inner_backend: + Optional pre-built backend for the hot tier. When *None* a + :class:`PlainKVCache` is created. + """ + + def __init__( + self, + num_layers: int, + num_heads: int, + head_dim: int, + *, + max_ram_tokens: int = 2048, + cache_dir: str | Path | None = None, + prompt_prefix: str = "", + inner_backend: KVCacheBackend | None = None, + ): + self.num_layers = num_layers + self.num_heads = num_heads + self.head_dim = head_dim + self.max_ram_tokens = max_ram_tokens + + if cache_dir is None: + cache_dir = Path.home() / ".cache" / "mlx-flash" / "kv-cache" + self._base_dir = Path(cache_dir) + self._session_dir = _session_dir(self._base_dir, prompt_prefix) + self._session_dir.mkdir(parents=True, exist_ok=True) + + self._hot: KVCacheBackend = inner_backend or PlainKVCache(num_layers, num_heads, head_dim) + + # Per-layer cold token count (only used for stats / merge decisions) + self._cold_tokens: list[int] = [0] * num_layers + + self._lock = threading.Lock() + + self._ram_hits = 0 + self._ssd_hits = 0 + self._ssd_writes = 0 + self._ssd_reads = 0 + self._bytes_on_ssd = 0 + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _ssd_path(self, layer_idx: int) -> Path: + return self._session_dir / f"layer_{layer_idx:04d}.safetensors" + + def _evict_to_ssd(self, layer_idx: int, keys: mx.array, values: mx.array) -> None: + """Append *keys*/*values* to the cold file for *layer_idx*.""" + path = self._ssd_path(layer_idx) + with self._lock: + if path.exists(): + existing = mx.load(str(path)) + old_k = existing["keys"] + old_v = existing["values"] + keys = mx.concatenate([old_k, keys], axis=1) + values = mx.concatenate([old_v, values], axis=1) + + mx.save_safetensors(str(path), {"keys": keys, "values": values}) + self._cold_tokens[layer_idx] = keys.shape[1] + self._ssd_writes += 1 + self._bytes_on_ssd = sum( + p.stat().st_size for p in self._session_dir.iterdir() if p.suffix == ".safetensors" + ) + + def _load_cold(self, layer_idx: int) -> tuple[mx.array, mx.array] | None: + path = self._ssd_path(layer_idx) + if not path.exists(): + return None + with self._lock: + data = mx.load(str(path)) + self._ssd_reads += 1 + self._ssd_hits += 1 + return data["keys"], data["values"] + + # ------------------------------------------------------------------ + # KVCacheBackend interface + # ------------------------------------------------------------------ + + def update(self, layer_idx: int, keys: mx.array, values: mx.array) -> tuple[mx.array, mx.array]: + all_k, all_v = self._hot.update(layer_idx, keys, values) + + seq_len = all_k.shape[1] + if seq_len > self.max_ram_tokens: + overflow = seq_len - self.max_ram_tokens + evict_k = all_k[:, :overflow, :] + evict_v = all_v[:, :overflow, :] + self._evict_to_ssd(layer_idx, evict_k, evict_v) + + keep_k = all_k[:, overflow:, :] + keep_v = all_v[:, overflow:, :] + + # Snapshot all layers before reset so we don't lose them + saved: list[tuple[mx.array, mx.array] | None] = [] + for i in range(self.num_layers): + if i == layer_idx: + saved.append((keep_k, keep_v)) + else: + ik, iv = self._hot.get_kv(i) + saved.append((ik, iv) if ik.shape[1] > 0 else None) + + self._hot.reset() + for i, pair in enumerate(saved): + if pair is not None: + self._hot.update(i, pair[0], pair[1]) + + all_k, all_v = keep_k, keep_v + + return all_k, all_v + + def get_kv(self, layer_idx: int) -> tuple[mx.array, mx.array]: + hot_k, hot_v = self._hot.get_kv(layer_idx) + self._ram_hits += 1 + + cold = self._load_cold(layer_idx) + if cold is None: + return hot_k, hot_v + + cold_k, cold_v = cold + if hot_k.shape[1] == 0: + return cold_k, cold_v + return ( + mx.concatenate([cold_k, hot_k], axis=1), + mx.concatenate([cold_v, hot_v], axis=1), + ) + + def reset(self) -> None: + self._hot.reset() + self._cold_tokens = [0] * self.num_layers + if self._session_dir.exists(): + shutil.rmtree(self._session_dir, ignore_errors=True) + self._session_dir.mkdir(parents=True, exist_ok=True) + self._ram_hits = 0 + self._ssd_hits = 0 + self._ssd_writes = 0 + self._ssd_reads = 0 + self._bytes_on_ssd = 0 + + def get_stats(self) -> dict[str, Any]: + return { + "strategy": "ssd", + "num_layers": self.num_layers, + "max_ram_tokens": self.max_ram_tokens, + "ram_hits": self._ram_hits, + "ssd_hits": self._ssd_hits, + "ssd_writes": self._ssd_writes, + "ssd_reads": self._ssd_reads, + "bytes_on_ssd": self._bytes_on_ssd, + "cold_tokens_per_layer": list(self._cold_tokens), + } diff --git a/mlx_flash_compress/telemetry.py b/mlx_flash_compress/telemetry.py new file mode 100644 index 0000000..a914948 --- /dev/null +++ b/mlx_flash_compress/telemetry.py @@ -0,0 +1,340 @@ +"""Hardware telemetry for Apple Silicon: GPU, power, thermal, and memory metrics. + +Uses IOReport-based sampling via ioreg and vm_stat to collect real-time +hardware metrics without requiring sudo or external dependencies. + +Usage: + from mlx_flash_compress.telemetry import HardwareTelemetry + + tel = HardwareTelemetry() + tel.start_sampling(interval_ms=1000) + sample = tel.sample() + stats = tel.get_stats() + tel.stop_sampling() +""" + +from __future__ import annotations + +import re +import subprocess +import threading +import time +from collections import deque +from dataclasses import asdict, dataclass, field +from typing import Optional + + +@dataclass +class TelemetrySample: + """A single point-in-time hardware telemetry snapshot.""" + + timestamp: float = 0.0 + gpu_util_pct: float = 0.0 + gpu_renderer_pct: float = 0.0 + gpu_tiler_pct: float = 0.0 + gpu_memory_used_gb: float = 0.0 + gpu_memory_total_gb: float = 0.0 + power_watts: float = 0.0 + cpu_temp_c: float = 0.0 + gpu_temp_c: float = 0.0 + memory_pressure: str = "normal" + memory_used_gb: float = 0.0 + memory_total_gb: float = 0.0 + ane_util_pct: float = 0.0 + + def to_dict(self) -> dict: + """Convert to JSON-serializable dict.""" + return asdict(self) + + +# Maximum number of samples kept in the circular buffer (2 minutes at 1s interval) +_MAX_HISTORY = 120 + + +class HardwareTelemetry: + """Thread-safe hardware telemetry collector for Apple Silicon. + + Samples GPU utilization, power draw, temperatures, and memory metrics + using macOS system commands (ioreg, vm_stat, sysctl). + """ + + def __init__(self) -> None: + self._history: deque[TelemetrySample] = deque(maxlen=_MAX_HISTORY) + self._lock = threading.Lock() + self._sampling = False + self._thread: Optional[threading.Thread] = None + self._total_memory_gb: float = self._detect_total_memory() + + # ── Public API ─────────────────────────────────────────────── + + def sample(self) -> TelemetrySample: + """Capture a single telemetry sample (current metrics).""" + s = TelemetrySample(timestamp=time.time()) + + # GPU metrics from IOAccelerator + self._read_gpu_metrics(s) + + # Power and thermal from SMC via ioreg + self._read_power_thermal(s) + + # System memory from vm_stat + sysctl + self._read_memory(s) + + with self._lock: + self._history.append(s) + + return s + + def start_sampling(self, interval_ms: int = 1000) -> None: + """Start continuous background sampling at the given interval.""" + if self._sampling: + return + self._sampling = True + self._thread = threading.Thread( + target=self._sampling_loop, + args=(interval_ms / 1000.0,), + daemon=True, + name="telemetry-sampler", + ) + self._thread.start() + + def stop_sampling(self) -> None: + """Stop the background sampling thread.""" + self._sampling = False + if self._thread is not None: + self._thread.join(timeout=5.0) + self._thread = None + + def get_history(self, seconds: int = 120) -> list[dict]: + """Return recent telemetry samples from the last N seconds.""" + cutoff = time.time() - seconds + with self._lock: + return [s.to_dict() for s in self._history if s.timestamp >= cutoff] + + def get_stats(self) -> dict: + """Return current values plus aggregates over the history window.""" + with self._lock: + samples = list(self._history) + + if not samples: + return { + "current": TelemetrySample(timestamp=time.time()).to_dict(), + "samples_count": 0, + "avg": {}, + "max": {}, + "min": {}, + } + + current = samples[-1] + count = len(samples) + + def _avg(attr: str) -> float: + vals = [getattr(s, attr) for s in samples] + return round(sum(vals) / len(vals), 2) if vals else 0.0 + + def _max(attr: str) -> float: + vals = [getattr(s, attr) for s in samples] + return round(max(vals), 2) if vals else 0.0 + + def _min(attr: str) -> float: + vals = [getattr(s, attr) for s in samples] + return round(min(vals), 2) if vals else 0.0 + + numeric_fields = [ + "gpu_util_pct", + "gpu_renderer_pct", + "gpu_tiler_pct", + "gpu_memory_used_gb", + "power_watts", + "cpu_temp_c", + "gpu_temp_c", + "memory_used_gb", + "ane_util_pct", + ] + + return { + "current": current.to_dict(), + "samples_count": count, + "avg": {f: _avg(f) for f in numeric_fields}, + "max": {f: _max(f) for f in numeric_fields}, + "min": {f: _min(f) for f in numeric_fields}, + } + + # ── Internal helpers ───────────────────────────────────────── + + def _sampling_loop(self, interval_s: float) -> None: + """Background sampling loop.""" + while self._sampling: + try: + self.sample() + except Exception: + pass # Never crash the sampling thread + time.sleep(interval_s) + + def _run_cmd(self, args: list[str], timeout: float = 3.0) -> str: + """Run a subprocess and return stdout, or empty string on failure.""" + try: + result = subprocess.run( + args, + capture_output=True, + text=True, + timeout=timeout, + ) + return result.stdout + except (subprocess.TimeoutExpired, FileNotFoundError, OSError): + return "" + + def _extract_ioreg_int(self, text: str, key: str) -> int: + """Extract an integer value from ioreg output by key name.""" + for line in text.splitlines(): + stripped = line.strip() + pos = stripped.find(key) + if pos >= 0: + after = stripped[pos + len(key) :] + digits = "" + started = False + for ch in after: + if ch.isdigit(): + digits += ch + started = True + elif started: + break + if digits: + return int(digits) + return 0 + + def _read_gpu_metrics(self, s: TelemetrySample) -> None: + """Read GPU utilization and memory from IOAccelerator.""" + text = self._run_cmd(["ioreg", "-r", "-d", "1", "-w", "0", "-c", "IOAccelerator"]) + if not text: + return + + s.gpu_util_pct = float(self._extract_ioreg_int(text, "Device Utilization %")) + s.gpu_renderer_pct = float(self._extract_ioreg_int(text, "Renderer Utilization %")) + s.gpu_tiler_pct = float(self._extract_ioreg_int(text, "Tiler Utilization %")) + + gpu_mem_used = self._extract_ioreg_int(text, '"In use system memory"=') + gpu_mem_alloc = self._extract_ioreg_int(text, '"Alloc system memory"=') + s.gpu_memory_used_gb = round(gpu_mem_used / (1024**3), 2) if gpu_mem_used else 0.0 + s.gpu_memory_total_gb = round(gpu_mem_alloc / (1024**3), 2) if gpu_mem_alloc else 0.0 + + def _read_power_thermal(self, s: TelemetrySample) -> None: + """Read power and temperature from AppleSMC via ioreg.""" + text = self._run_cmd(["ioreg", "-r", "-n", "AppleSmartBattery"]) + if text: + # Try to extract instantaneous power from battery info + watts = self._extract_ioreg_int(text, '"InstantAmperage"=') + voltage = self._extract_ioreg_int(text, '"Voltage"=') + if watts and voltage: + # Convert mA * mV -> W (both are in milli units) + s.power_watts = round(abs(watts) * voltage / 1_000_000, 1) + + # Try to get temperatures from thermal sensors + thermal_text = self._run_cmd(["ioreg", "-r", "-d", "1", "-w", "0", "-c", "AppleARMIODevice"]) + if thermal_text: + # Look for CPU/GPU temperature entries + cpu_temp = self._extract_temperature(thermal_text, "cpu") + gpu_temp = self._extract_temperature(thermal_text, "gpu") + if cpu_temp > 0: + s.cpu_temp_c = cpu_temp + if gpu_temp > 0: + s.gpu_temp_c = gpu_temp + + # Fallback: try sysctl for CPU temperature (macOS 14+) + if s.cpu_temp_c == 0.0: + temp_text = self._run_cmd(["sysctl", "-n", "machdep.xcpm.cpu_thermal_level"]) + if temp_text.strip(): + try: + level = int(temp_text.strip()) + # Thermal level is 0-100 scale, approximate to Celsius + s.cpu_temp_c = round(35.0 + level * 0.6, 1) + except ValueError: + pass + + def _extract_temperature(self, text: str, sensor_type: str) -> float: + """Extract temperature value from ioreg thermal sensor output.""" + # Look for temperature patterns like "temperature" = XX or sensor readings + pattern = re.compile( + rf'"(?:.*{sensor_type}.*temp|{sensor_type}.*die).*?".*?=\s*(\d+)', + re.IGNORECASE, + ) + match = pattern.search(text) + if match: + raw = int(match.group(1)) + # Values might be in centi-degrees or direct Celsius + if raw > 200: + return round(raw / 100.0, 1) + elif raw > 0: + return float(raw) + return 0.0 + + def _read_memory(self, s: TelemetrySample) -> None: + """Read system memory stats from vm_stat and sysctl.""" + s.memory_total_gb = self._total_memory_gb + + # Try psutil first (most accurate), fall back to vm_stat + try: + import psutil + + vm = psutil.virtual_memory() + s.memory_used_gb = round((vm.total - vm.available) / (1024**3), 2) + s.memory_total_gb = round(vm.total / (1024**3), 2) + pct = vm.percent + if pct > 90: + s.memory_pressure = "critical" + elif pct > 75: + s.memory_pressure = "warn" + else: + s.memory_pressure = "normal" + return + except ImportError: + pass + + # Fallback: parse vm_stat + text = self._run_cmd(["vm_stat"]) + if not text: + return + + page_size = 16384 # Default for Apple Silicon + ps_match = re.search(r"page size of (\d+) bytes", text) + if ps_match: + page_size = int(ps_match.group(1)) + + def _pages(key: str) -> int: + m = re.search(rf"{key}:\s+(\d+)", text) + return int(m.group(1)) if m else 0 + + free = _pages("Pages free") + active = _pages("Pages active") + inactive = _pages("Pages inactive") + wired = _pages("Pages wired down") + compressed = _pages("Pages occupied by compressor") + + used_pages = active + wired + compressed + total_pages = free + active + inactive + wired + compressed + if total_pages > 0: + s.memory_used_gb = round(used_pages * page_size / (1024**3), 2) + used_pct = used_pages / total_pages * 100 + if used_pct > 90: + s.memory_pressure = "critical" + elif used_pct > 75: + s.memory_pressure = "warn" + else: + s.memory_pressure = "normal" + + # Memory pressure from macOS memory_pressure command + pressure_text = self._run_cmd(["memory_pressure"]) + if "critical" in pressure_text.lower(): + s.memory_pressure = "critical" + elif "warn" in pressure_text.lower(): + s.memory_pressure = "warn" + + def _detect_total_memory(self) -> float: + """Detect total system memory via sysctl.""" + text = self._run_cmd(["sysctl", "-n", "hw.memsize"]) + if text.strip(): + try: + return round(int(text.strip()) / (1024**3), 2) + except ValueError: + pass + return 0.0 diff --git a/scripts/train_dflash_drafter.py b/scripts/train_dflash_drafter.py index bdacec0..52b9cd6 100644 --- a/scripts/train_dflash_drafter.py +++ b/scripts/train_dflash_drafter.py @@ -78,6 +78,18 @@ def select_checkpoint_layers(num_layers: int, n: int = 5) -> list[int]: "recommended_lr": 3e-4, "notes": "284B total / 13B active MoE. Use 2-bit on 64GB, 4-bit on 128GB+.", }, + "deepseek-v4-flash-mini": { + "target_2bit": "mlx-community/DeepSeek-V4-Flash-2bit-DQ", + "target_4bit": "mlx-community/DeepSeek-V4-Flash-4bit", + "drafter_layers": 3, + "drafter_hidden": 1024, + "block_size": 16, + "recommended_steps": 8000, + "recommended_samples": 800, + "recommended_lr": 4e-4, + "notes": "Mini drafter (3 layers, 1024 hidden) for RAM-constrained setups. " + "Trades ~5-10% acceptance for ~60% less drafter memory.", + }, "qwen3.6-35b-a3b": { "target_4bit": "mlx-community/Qwen3.6-35B-A3B-4bit", "drafter_hub": "z-lab/Qwen3.6-35B-A3B-DFlash", @@ -92,6 +104,64 @@ def select_checkpoint_layers(num_layers: int, n: int = 5) -> list[int]: } +def _get_available_ram_gb() -> float: + """Return available system RAM in GB (best-effort).""" + try: + import subprocess + + result = subprocess.run( + ["sysctl", "-n", "hw.memsize"], + capture_output=True, + text=True, + ) + if result.returncode == 0: + return int(result.stdout.strip()) / (1024**3) + except Exception: + pass + # Fallback: try psutil + try: + import psutil + + return psutil.virtual_memory().total / (1024**3) + except Exception: + pass + return 0.0 + + +def auto_select_preset(target_model: str | None = None) -> str: + """Automatically select the best drafter preset based on available RAM. + + Heuristic: + - If target is DeepSeek V4 Flash: + - >= 96 GB RAM: deepseek-v4-flash (full drafter) + - < 96 GB RAM: deepseek-v4-flash-mini (smaller drafter) + - Otherwise defaults to qwen3.6-35b-a3b + + Args: + target_model: Optional target model path/name to help select. + + Returns: + Preset name string. + """ + ram_gb = _get_available_ram_gb() + + # Detect DeepSeek V4 Flash from target model name + is_deepseek_v4 = False + if target_model: + lower = target_model.lower() + if "deepseek" in lower and ("v4" in lower or "flash" in lower): + is_deepseek_v4 = True + + if is_deepseek_v4: + if ram_gb >= 96: + return "deepseek-v4-flash" + else: + return "deepseek-v4-flash-mini" + + # Default + return "qwen3.6-35b-a3b" + + def detect_target(model_path: str): """Load target model and detect architecture.""" print(f"Loading target model: {model_path}") @@ -185,45 +255,76 @@ def collect_hidden_states(args): } (output_dir / "model_info.json").write_text(json.dumps(model_info, indent=2)) - # Diverse prompts covering code, math, reasoning, prose, dialogue - prompts = [ - # Code (Python, JS, Rust, SQL, shell) - "def fibonacci(n):\n ", - "import torch\nimport torch.nn as nn\n\nclass", - "async def fetch_data(url: str) -> dict:\n ", - "def merge_sort(arr):\n if len(arr) <= 1:\n return arr\n", - "class BinaryTree:\n def __init__(self, value):\n", - "SELECT u.name, COUNT(o.id) FROM users u JOIN", - "function debounce(fn, delay) {\n let timer;\n return function(", - "fn main() -> Result<(), Box> {\n let client =", - "#!/bin/bash\nset -euo pipefail\n\nfor file in", - "def train_step(model, optimizer, batch):\n optimizer.zero_grad()\n", - "class LRUCache:\n def __init__(self, capacity: int):\n", - "@app.route('/api/users', methods=['POST'])\ndef create_user():\n", - # Math & reasoning - "To solve the quadratic equation ax^2 + bx + c = 0,", - "Given a matrix A of size m x n, the transpose A^T is", - "The gradient descent algorithm works by iteratively", - "Prove by induction that the sum of the first n natural numbers is", - "The eigenvalues of a 2x2 matrix can be found by solving", - "Using the chain rule, the derivative of f(g(x)) is", - # Technical prose - "The transformer architecture consists of an encoder and decoder, where", - "The attention mechanism allows the model to focus on", - "In Python, list comprehensions provide a concise way to", - "The CAP theorem states that a distributed system cannot simultaneously", - "Garbage collection in modern languages uses either reference counting or", - "The difference between TCP and UDP is that TCP provides", - # Dialogue & instruction following - "User: How do I implement a binary search?\nAssistant:", - "Explain the difference between a stack and a queue to", - "Write a Python function that takes a list of integers and returns", - "Summarize the key advantages of using a hash table for", - # Structured output - '{"name": "John", "age": 30, "skills": [', - "# README\n\n## Installation\n\n```bash\npip install", - "| Column A | Column B | Column C |\n|----------|----------|----------|\n|", - ] + # Load calibration prompts from file if provided, otherwise use built-in defaults + if getattr(args, "calibration_prompts", None): + prompts_path = Path(args.calibration_prompts) + if not prompts_path.exists(): + print(f"Error: calibration prompts file not found: {prompts_path}") + sys.exit(1) + prompts = [] + with open(prompts_path) as f: + for line in f: + line = line.strip() + if not line: + continue + try: + data = json.loads(line) + # Support {"prompt": "..."} or {"text": "..."} or plain string + if isinstance(data, dict): + prompt_text = data.get("prompt") or data.get("text") or "" + elif isinstance(data, str): + prompt_text = data + else: + continue + if prompt_text: + prompts.append(prompt_text) + except json.JSONDecodeError: + # Treat as plain text line + prompts.append(line) + print(f" Loaded {len(prompts)} calibration prompts from {prompts_path}") + if not prompts: + print("Error: no valid prompts found in calibration file") + sys.exit(1) + else: + # Diverse prompts covering code, math, reasoning, prose, dialogue + prompts = [ + # Code (Python, JS, Rust, SQL, shell) + "def fibonacci(n):\n ", + "import torch\nimport torch.nn as nn\n\nclass", + "async def fetch_data(url: str) -> dict:\n ", + "def merge_sort(arr):\n if len(arr) <= 1:\n return arr\n", + "class BinaryTree:\n def __init__(self, value):\n", + "SELECT u.name, COUNT(o.id) FROM users u JOIN", + "function debounce(fn, delay) {\n let timer;\n return function(", + "fn main() -> Result<(), Box> {\n let client =", + "#!/bin/bash\nset -euo pipefail\n\nfor file in", + "def train_step(model, optimizer, batch):\n optimizer.zero_grad()\n", + "class LRUCache:\n def __init__(self, capacity: int):\n", + "@app.route('/api/users', methods=['POST'])\ndef create_user():\n", + # Math & reasoning + "To solve the quadratic equation ax^2 + bx + c = 0,", + "Given a matrix A of size m x n, the transpose A^T is", + "The gradient descent algorithm works by iteratively", + "Prove by induction that the sum of the first n natural numbers is", + "The eigenvalues of a 2x2 matrix can be found by solving", + "Using the chain rule, the derivative of f(g(x)) is", + # Technical prose + "The transformer architecture consists of an encoder and decoder, where", + "The attention mechanism allows the model to focus on", + "In Python, list comprehensions provide a concise way to", + "The CAP theorem states that a distributed system cannot simultaneously", + "Garbage collection in modern languages uses either reference counting or", + "The difference between TCP and UDP is that TCP provides", + # Dialogue & instruction following + "User: How do I implement a binary search?\nAssistant:", + "Explain the difference between a stack and a queue to", + "Write a Python function that takes a list of integers and returns", + "Summarize the key advantages of using a hash table for", + # Structured output + '{"name": "John", "age": 30, "skills": [', + "# README\n\n## Installation\n\n```bash\npip install", + "| Column A | Column B | Column C |\n|----------|----------|----------|\n|", + ] from mlx_lm import generate @@ -533,9 +634,22 @@ def show_preset(args): for name, cfg in DRAFTER_PRESETS.items(): print(f"\n {name}:") print(f" {cfg['notes']}") + print("\n auto:") + print(" Automatically selects preset based on available RAM and target model.") print("\nUsage: python scripts/train_dflash_drafter.py preset ") return + if args.name == "auto": + target = getattr(args, "target_model", None) + selected = auto_select_preset(target) + ram_gb = _get_available_ram_gb() + print(f"Auto-selected preset: {selected}") + print(f" Available RAM: {ram_gb:.0f} GB") + if target: + print(f" Target model: {target}") + print(f"\nShowing configuration for '{selected}':\n") + args.name = selected + cfg = DRAFTER_PRESETS[args.name] target_key = f"target_{args.quant}" target = cfg.get(target_key, cfg.get("target_4bit", "UNKNOWN")) @@ -599,6 +713,13 @@ def main(): collect_p.add_argument( "--checkpoint-layers", type=str, default=None, help="Comma-separated checkpoint layer IDs (e.g. 1,10,19,28,37)" ) + collect_p.add_argument( + "--calibration-prompts", + type=str, + default=None, + help="Path to a .jsonl file of custom prompts for calibration data collection. " + 'Each line: {"prompt": "..."} or plain text. If omitted, uses built-in diverse prompts.', + ) train_p = subparsers.add_parser("train", help="Train DFlash drafter") train_p.add_argument("--training-data", type=str, required=True) @@ -625,8 +746,14 @@ def main(): type=str, nargs="?", default=None, - choices=list(DRAFTER_PRESETS.keys()), - help="Preset name (omit to list all)", + choices=list(DRAFTER_PRESETS.keys()) + ["auto"], + help="Preset name (omit to list all, 'auto' to auto-select based on RAM)", + ) + preset_p.add_argument( + "--target-model", + type=str, + default=None, + help="Target model name (used by 'auto' to pick the best preset)", ) preset_p.add_argument("--quant", type=str, default="4bit", choices=["2bit", "4bit"], help="Quantization level") diff --git a/tests/test_dflash_acceptance.py b/tests/test_dflash_acceptance.py new file mode 100644 index 0000000..a76bbe2 --- /dev/null +++ b/tests/test_dflash_acceptance.py @@ -0,0 +1,359 @@ +"""Tests for DFlash acceptance improvements: soft threshold, hidden state +normalization, multi-candidate top-K verification, and block size auto-tuning. + +These tests focus on the new DFlashRunner features without requiring +actual model downloads — all use mock/fake models. +""" + +from __future__ import annotations + +import pytest + +try: + import mlx.core as mx + import mlx.nn as nn + + MLX_AVAILABLE = True +except ImportError: + MLX_AVAILABLE = False + +from mlx_flash_compress.dflash_model import ( + DFlashDraftModel, + DFlashModelConfig, + DFlashRunner, +) + +pytestmark = pytest.mark.skipif(not MLX_AVAILABLE, reason="MLX not available") + + +# -- Helpers -- + + +class FakeLayer(nn.Module): + def __init__(self, dim): + super().__init__() + self.linear = nn.Linear(dim, dim, bias=False) + + def __call__(self, x, **kwargs): + return self.linear(x) + + +class FakeModel(nn.Module): + def __init__(self, num_layers=10, hidden_size=64, vocab_size=100): + super().__init__() + self.model = _FakeInner(num_layers, hidden_size, vocab_size) + self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False) + + def __call__(self, input_ids, **kwargs): + h = self.model.embed_tokens(input_ids) + for layer in self.model.layers: + h = layer(h) + h = self.model.norm(h) + return self.lm_head(h) + + +class _FakeInner(nn.Module): + def __init__(self, num_layers, hidden_size, vocab_size): + super().__init__() + self.embed_tokens = nn.Embedding(vocab_size, hidden_size) + self.layers = [FakeLayer(hidden_size) for _ in range(num_layers)] + self.norm = nn.RMSNorm(hidden_size) + + +class FakeTokenizer: + eos_token_id = 2 + + def __init__(self, vocab_size=100): + self._vocab_size = vocab_size + + def encode(self, text): + return [1] + [ord(c) % self._vocab_size for c in text[:20]] + + def decode(self, ids): + return "".join(chr(max(32, i % 128)) for i in ids) + + +def _make_runner( + num_layers=10, + hidden_size=64, + vocab_size=100, + block_size=4, + **kwargs, +): + target = FakeModel(num_layers=num_layers, hidden_size=hidden_size, vocab_size=vocab_size) + mx.eval(target.parameters()) + tokenizer = FakeTokenizer(vocab_size) + + config = DFlashModelConfig( + hidden_size=hidden_size, + intermediate_size=hidden_size * 2, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=hidden_size // 4, + block_size=block_size, + vocab_size=vocab_size, + target_layer_ids=[1, 4, 7], + mask_token_id=0, + ) + drafter = DFlashDraftModel(config) + mx.eval(drafter.parameters()) + + runner = DFlashRunner(target, tokenizer, drafter, config, **kwargs) + return runner + + +# ============================================================ +# Soft acceptance (threshold-based) +# ============================================================ + + +class TestSoftAcceptance: + def test_threshold_zero_is_strict_greedy(self): + """acceptance_threshold=0 should behave identically to strict greedy.""" + runner = _make_runner(acceptance_threshold=0.0) + text, stats = runner.generate("hello", max_tokens=8, use_cache=False) + assert isinstance(text, str) + assert 0 <= stats["acceptance_rate"] <= 1.0 + + def test_threshold_positive_accepts_more(self): + """A positive threshold should generally accept >= what greedy accepts.""" + runner_strict = _make_runner(acceptance_threshold=0.0) + runner_soft = _make_runner(acceptance_threshold=0.01) + + # Run both on same prompt + _, stats_strict = runner_strict.generate("hello world", max_tokens=8, use_cache=False) + _, stats_soft = runner_soft.generate("hello world", max_tokens=8, use_cache=False) + + # Soft should accept at least as much (or more) than strict + # Not guaranteed to be strictly >= due to randomness in fake models, + # but both should produce valid results + assert stats_strict["tokens_generated"] > 0 + assert stats_soft["tokens_generated"] > 0 + + def test_threshold_via_generate_kwarg(self): + """acceptance_threshold can be passed to generate() directly.""" + runner = _make_runner() + text, stats = runner.generate( + "test prompt", + max_tokens=6, + use_cache=False, + acceptance_threshold=0.05, + ) + assert isinstance(text, str) + assert stats["tokens_generated"] > 0 + + def test_threshold_cached_mode(self): + """Soft acceptance should work with cached generation too.""" + runner = _make_runner(acceptance_threshold=0.01) + text, stats = runner.generate("cached test", max_tokens=6, use_cache=True) + assert isinstance(text, str) + assert stats["tokens_generated"] > 0 + + def test_verify_with_threshold_unit(self): + """Unit test _verify_with_threshold with controlled logits.""" + runner = _make_runner() + + # Create predictions where token 5 has high probability + # Shape: [1, 3, vocab_size] + vocab = 100 + logits = mx.zeros((1, 3, vocab)) + # Position 0: token 5 gets high logit -> high prob + logits = logits.at[0, 0, 5].add(10.0) + # Position 1: token 7 gets moderate logit + logits = logits.at[0, 1, 7].add(5.0) + # Position 2: token 3 gets low logit + logits = logits.at[0, 2, 3].add(0.5) + + # Draft tokens: [5, 7, 3] + draft = [5, 7, 3] + + # With threshold=0, should use greedy (all match -> 3 accepted) + n = runner._verify_with_threshold(logits, draft, accept_top_k=1, acceptance_threshold=0.0) + assert n == 3 # greedy: argmax at each position matches draft + + # With threshold=0.5, token at position 2 might not pass + # because its probability after softmax of 0.5 logit is low + n_soft = runner._verify_with_threshold(logits, draft, accept_top_k=1, acceptance_threshold=0.5) + # First two should pass (high prob), third might not + assert n_soft >= 2 + + +# ============================================================ +# Hidden state normalization +# ============================================================ + + +class TestHiddenStateNormalization: + def test_normalization_enabled(self): + """normalize_hidden=True should produce valid generation.""" + runner = _make_runner(normalize_hidden=True) + text, stats = runner.generate("normalize test", max_tokens=6, use_cache=False) + assert isinstance(text, str) + assert stats["tokens_generated"] > 0 + + def test_normalization_disabled_default(self): + """Default should be normalize_hidden=False.""" + runner = _make_runner() + assert runner._normalize_hidden is False + + def test_apply_hidden_normalization_unit(self): + """Unit test the static normalization method.""" + h = mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + normalized = DFlashRunner._apply_hidden_normalization(h) + mx.eval(normalized) + + # Each row should have approximately unit norm + norms = mx.sqrt(mx.sum(normalized * normalized, axis=-1)) + mx.eval(norms) + for i in range(norms.shape[0]): + assert abs(float(norms[i].item()) - 1.0) < 1e-4 + + def test_normalization_cached_mode(self): + """normalize_hidden should work with cached generation.""" + runner = _make_runner(normalize_hidden=True) + text, stats = runner.generate("cached normalize", max_tokens=6, use_cache=True) + assert isinstance(text, str) + assert stats["tokens_generated"] > 0 + + +# ============================================================ +# Multi-candidate top-K verification +# ============================================================ + + +class TestTopKVerification: + def test_top_k_1_is_greedy(self): + """accept_top_k=1 should be equivalent to greedy.""" + runner = _make_runner() + text, stats = runner.generate("top k test", max_tokens=6, use_cache=False, accept_top_k=1) + assert isinstance(text, str) + assert stats["tokens_generated"] > 0 + + def test_top_k_5_accepts_more(self): + """accept_top_k=5 should accept at least as many as top_k=1.""" + runner_k1 = _make_runner() + runner_k5 = _make_runner() + + _, stats_k1 = runner_k1.generate("test", max_tokens=8, use_cache=False, accept_top_k=1) + _, stats_k5 = runner_k5.generate("test", max_tokens=8, use_cache=False, accept_top_k=5) + + # Both should produce valid output + assert stats_k1["tokens_generated"] > 0 + assert stats_k5["tokens_generated"] > 0 + + def test_top_k_cached(self): + """Top-K verification in cached mode.""" + runner = _make_runner() + text, stats = runner.generate("cached top k", max_tokens=6, use_cache=True, accept_top_k=3) + assert isinstance(text, str) + assert stats["tokens_generated"] > 0 + + def test_verify_with_threshold_top_k_unit(self): + """Unit test top-K verification logic.""" + runner = _make_runner() + + # Create predictions: position 0 has token 5 as #1 and token 10 as #2 + vocab = 100 + logits = mx.zeros((1, 2, vocab)) + logits = logits.at[0, 0, 5].add(10.0) + logits = logits.at[0, 0, 10].add(9.0) + logits = logits.at[0, 1, 7].add(10.0) + + # Draft [10, 7]: token 10 is #2 at position 0 + # With top_k=1 (greedy): 10 != 5 -> 0 accepted + n_k1 = runner._verify_with_threshold(logits, [10, 7], accept_top_k=1, acceptance_threshold=0.0) + assert n_k1 == 0 # 10 is not argmax (5 is) + + # With top_k=2: 10 is in top-2 -> accepted, then 7 is top-1 -> accepted + n_k2 = runner._verify_with_threshold(logits, [10, 7], accept_top_k=2, acceptance_threshold=0.0) + assert n_k2 == 2 + + +# ============================================================ +# Block size auto-tuning +# ============================================================ + + +class TestBlockSizeAutoTuning: + def test_auto_block_size_disabled_by_default(self): + """auto_block_size should default to False.""" + runner = _make_runner() + assert runner._auto_block_size is False + # block_size should not change during generation + initial_bs = runner.block_size + runner._auto_tune_block_size(num_accepted=5, n_draft=10) + assert runner.block_size == initial_bs + + def test_auto_tune_increases_on_high_acceptance(self): + """Block size should increase after 10 rounds with >40% acceptance.""" + runner = _make_runner(auto_block_size=True, block_size=4) + assert runner.block_size == 4 + + # Simulate 10 rounds with 50% acceptance (> 40% threshold) + for _ in range(10): + runner._auto_tune_block_size(num_accepted=5, n_draft=10) + + # Should have doubled to 8 + assert runner.block_size == 8 + + def test_auto_tune_decreases_on_low_acceptance(self): + """Block size should decrease after 10 rounds with <15% acceptance.""" + runner = _make_runner(auto_block_size=True, block_size=4) + + # Simulate 10 rounds with 10% acceptance (< 15% threshold) + for _ in range(10): + runner._auto_tune_block_size(num_accepted=1, n_draft=10) + + # Should have halved to 2 + assert runner.block_size == 2 + + def test_auto_tune_no_change_in_middle_range(self): + """Block size should not change when acceptance is between 15-40%.""" + runner = _make_runner(auto_block_size=True, block_size=4) + + # Simulate 10 rounds with 25% acceptance (between 15-40%) + for _ in range(10): + runner._auto_tune_block_size(num_accepted=25, n_draft=100) + + # Should stay at 4 + assert runner.block_size == 4 + + def test_auto_tune_clamps_at_minimum(self): + """Block size should not go below 2.""" + runner = _make_runner(auto_block_size=True, block_size=2) + + # Simulate 10 rounds with very low acceptance + for _ in range(10): + runner._auto_tune_block_size(num_accepted=0, n_draft=10) + + assert runner.block_size == 2 # minimum + + def test_auto_tune_clamps_at_maximum(self): + """Block size should not go above 16.""" + runner = _make_runner(auto_block_size=True, block_size=16) + + # Simulate 10 rounds with high acceptance + for _ in range(10): + runner._auto_tune_block_size(num_accepted=8, n_draft=10) + + assert runner.block_size == 16 # maximum, no change + + def test_auto_tune_resets_counters(self): + """After 10 rounds, counters should reset for the next window.""" + runner = _make_runner(auto_block_size=True, block_size=4) + + for _ in range(10): + runner._auto_tune_block_size(num_accepted=5, n_draft=10) + + # After first window: block_size = 8, counters reset + assert runner._auto_bs_rounds == 0 + assert runner._auto_bs_accepted == 0 + assert runner._auto_bs_drafted == 0 + + def test_auto_tune_in_generation(self): + """Auto-tune should work during actual generation.""" + runner = _make_runner(auto_block_size=True, block_size=4) + text, stats = runner.generate("auto tune test", max_tokens=8, use_cache=False) + assert isinstance(text, str) + assert stats["tokens_generated"] > 0 diff --git a/tests/test_draft_expert_prefetch.py b/tests/test_draft_expert_prefetch.py new file mode 100644 index 0000000..5739cac --- /dev/null +++ b/tests/test_draft_expert_prefetch.py @@ -0,0 +1,201 @@ +"""Tests for DraftExpertPrefetcher — expert prefetch from draft predictions.""" + +from __future__ import annotations + +import pytest + +try: + import mlx.core as mx + import mlx.nn as nn + + HAS_MLX = True +except (ImportError, ModuleNotFoundError): + HAS_MLX = False + +pytestmark = pytest.mark.skipif(not HAS_MLX, reason="requires mlx") + + +# -- Mock objects -- + + +class MockExpertCache: + """Mock cache that records prefetch calls.""" + + def __init__(self): + self.prefetch_calls: list[tuple[int, list[int]]] = [] + + def prefetch(self, layer_idx: int, expert_ids: list[int]): + self.prefetch_calls.append((layer_idx, expert_ids)) + + +class MockEmbed(nn.Module): + """Minimal embedding function for testing.""" + + def __init__(self, vocab_size: int = 100, hidden_size: int = 32): + super().__init__() + self.embed = nn.Embedding(vocab_size, hidden_size) + + def __call__(self, x): + return self.embed(x) + + +# -- Tests -- + + +class TestDraftExpertPrefetcher: + def _make_prefetcher(self, num_layers=3, num_experts=8, hidden_size=32, top_k=2): + from mlx_flash_compress.draft_expert_prefetch import DraftExpertPrefetcher + + cache = MockExpertCache() + # Create router weights: layer_idx -> [num_experts, hidden_size] + router_weights = {} + for layer_idx in range(num_layers): + w = mx.random.normal((num_experts, hidden_size)) + mx.eval(w) + router_weights[layer_idx] = w + + embed = MockEmbed(vocab_size=100, hidden_size=hidden_size) + mx.eval(embed.parameters()) + + prefetcher = DraftExpertPrefetcher( + expert_cache=cache, + router_weights=router_weights, + top_k=top_k, + ) + return prefetcher, cache, embed + + def test_prefetch_returns_layer_expert_pairs(self): + prefetcher, cache, embed = self._make_prefetcher( + num_layers=3, + num_experts=8, + top_k=2, + ) + draft_ids = [10, 20, 30] + pairs = prefetcher.prefetch_from_drafts(draft_ids, embed) + + # Should return top_k experts per layer + assert len(pairs) == 3 * 2 # 3 layers x 2 top_k + for layer_idx, expert_idx in pairs: + assert 0 <= layer_idx < 3 + assert 0 <= expert_idx < 8 + + def test_prefetch_triggers_cache_calls(self): + prefetcher, cache, embed = self._make_prefetcher( + num_layers=2, + num_experts=4, + top_k=2, + ) + draft_ids = [5, 15] + prefetcher.prefetch_from_drafts(draft_ids, embed) + + # Should have called prefetch on the cache for each layer + assert len(cache.prefetch_calls) == 2 + for layer_idx, expert_ids in cache.prefetch_calls: + assert 0 <= layer_idx < 2 + assert len(expert_ids) == 2 + + def test_empty_draft_ids(self): + prefetcher, cache, embed = self._make_prefetcher() + pairs = prefetcher.prefetch_from_drafts([], embed) + assert pairs == [] + assert len(cache.prefetch_calls) == 0 + + def test_empty_router_weights(self): + from mlx_flash_compress.draft_expert_prefetch import DraftExpertPrefetcher + + cache = MockExpertCache() + embed = MockEmbed() + mx.eval(embed.parameters()) + + prefetcher = DraftExpertPrefetcher( + expert_cache=cache, + router_weights={}, + top_k=2, + ) + pairs = prefetcher.prefetch_from_drafts([1, 2, 3], embed) + assert pairs == [] + + def test_stats_tracking(self): + prefetcher, cache, embed = self._make_prefetcher( + num_layers=2, + num_experts=4, + top_k=2, + ) + stats = prefetcher.get_stats() + assert stats["prefetch_requests"] == 0 + assert stats["total_predictions"] == 0 + assert stats["accuracy"] == 0.0 + + prefetcher.prefetch_from_drafts([5, 10], embed) + + stats = prefetcher.get_stats() + assert stats["prefetch_requests"] == 2 # one per layer + assert stats["total_predictions"] == 4 # 2 layers x 2 top_k + + def test_record_cache_hit(self): + prefetcher, cache, embed = self._make_prefetcher( + num_layers=1, + num_experts=4, + top_k=2, + ) + prefetcher.prefetch_from_drafts([5], embed) + prefetcher.record_cache_hit(count=2) + + stats = prefetcher.get_stats() + assert stats["cache_hits_from_prefetch"] == 2 + assert stats["accuracy"] == 2 / 2 # 2 hits / 2 predictions + + def test_none_cache_graceful(self): + """Prefetcher with None cache should work (no-op prefetch).""" + from mlx_flash_compress.draft_expert_prefetch import DraftExpertPrefetcher + + embed = MockEmbed() + mx.eval(embed.parameters()) + + router_weights = {0: mx.random.normal((4, 32))} + mx.eval(router_weights[0]) + + prefetcher = DraftExpertPrefetcher( + expert_cache=None, + router_weights=router_weights, + top_k=2, + ) + pairs = prefetcher.prefetch_from_drafts([1, 2, 3], embed) + + # Should still return predictions even without a cache + assert len(pairs) == 2 # 1 layer x 2 top_k + stats = prefetcher.get_stats() + # No prefetch_requests because cache is None (no prefetch call made) + assert stats["prefetch_requests"] == 0 + assert stats["total_predictions"] == 2 + + def test_router_weight_transposed_layout(self): + """Router weights in [hidden_size, num_experts] layout should also work.""" + from mlx_flash_compress.draft_expert_prefetch import DraftExpertPrefetcher + + cache = MockExpertCache() + embed = MockEmbed(hidden_size=32) + mx.eval(embed.parameters()) + + # [hidden_size, num_experts] layout (transposed from typical) + router_weights = {0: mx.random.normal((32, 8))} + mx.eval(router_weights[0]) + + prefetcher = DraftExpertPrefetcher( + expert_cache=cache, + router_weights=router_weights, + top_k=3, + ) + pairs = prefetcher.prefetch_from_drafts([1, 2], embed) + assert len(pairs) == 3 # 1 layer x 3 top_k + + def test_top_k_clamped_to_num_experts(self): + """top_k larger than num_experts should be clamped.""" + prefetcher, cache, embed = self._make_prefetcher( + num_layers=1, + num_experts=3, + top_k=10, + ) + pairs = prefetcher.prefetch_from_drafts([1, 2], embed) + # Should get at most num_experts (3) results + assert len(pairs) == 3 diff --git a/tests/test_imports.py b/tests/test_imports.py index 2fa5124..7caf890 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -86,6 +86,13 @@ "HybridKVCache": ("mlx_flash_compress.kv_cache_backend", "HybridKVCache"), "create_kv_cache": ("mlx_flash_compress.kv_cache_backend", "create_kv_cache"), "install_kv_cache": ("mlx_flash_compress.kv_cache_backend", "install_kv_cache"), + # ssd_kv_cache.py + "SSDKVCache": ("mlx_flash_compress.ssd_kv_cache", "SSDKVCache"), + # prefix_cache.py + "PrefixCacheTrie": ("mlx_flash_compress.prefix_cache", "PrefixCacheTrie"), + "get_or_compute_prefix": ("mlx_flash_compress.prefix_cache", "get_or_compute_prefix"), + # draft_expert_prefetch.py + "DraftExpertPrefetcher": ("mlx_flash_compress.draft_expert_prefetch", "DraftExpertPrefetcher"), } diff --git a/tests/test_prefix_cache.py b/tests/test_prefix_cache.py new file mode 100644 index 0000000..09ecab0 --- /dev/null +++ b/tests/test_prefix_cache.py @@ -0,0 +1,325 @@ +"""Tests for RadixAttention-style prefix caching trie.""" + +from __future__ import annotations + +import threading +import time + +import pytest + +try: + import mlx.core as mx + + HAS_MLX = True +except ImportError: + HAS_MLX = False + +from mlx_flash_compress.prefix_cache import ( + PrefixCacheTrie, + TrieNode, + get_or_compute_prefix, +) + +pytestmark = pytest.mark.skipif(not HAS_MLX, reason="MLX required") + +NUM_LAYERS = 4 +NUM_HEADS = 2 +HEAD_DIM = 64 + + +def _make_kv_list(seq_len, num_layers=NUM_LAYERS, seed=0): + mx.random.seed(seed) + result = [] + for _ in range(num_layers): + k = mx.random.normal((NUM_HEADS, seq_len, HEAD_DIM)) + v = mx.random.normal((NUM_HEADS, seq_len, HEAD_DIM)) + result.append((k, v)) + return result + + +class TestInsertAndLookup: + def test_insert_and_exact_lookup(self): + trie = PrefixCacheTrie() + tokens = [1, 2, 3, 4, 5] + kv = _make_kv_list(5) + trie.insert(tokens, kv) + + length, result = trie.lookup(tokens) + assert length == 5 + assert result is not None + assert len(result) == NUM_LAYERS + + def test_lookup_no_match(self): + trie = PrefixCacheTrie() + trie.insert([1, 2, 3], _make_kv_list(3)) + + length, result = trie.lookup([9, 8, 7]) + assert length == 0 + assert result is None + + def test_lookup_empty_trie(self): + trie = PrefixCacheTrie() + length, result = trie.lookup([1, 2, 3]) + assert length == 0 + assert result is None + + def test_insert_empty_tokens(self): + trie = PrefixCacheTrie() + kv = _make_kv_list(0) + trie.insert([], kv) + length, result = trie.lookup([]) + assert length == 0 + + +class TestPrefixMatching: + def test_longest_prefix_match(self): + trie = PrefixCacheTrie() + kv3 = _make_kv_list(3, seed=1) + kv5 = _make_kv_list(5, seed=2) + trie.insert([1, 2, 3], kv3) + trie.insert([1, 2, 3, 4, 5], kv5) + + length, result = trie.lookup([1, 2, 3, 4, 5, 6, 7]) + assert length == 5 + assert result is kv5 + + def test_partial_prefix_match(self): + trie = PrefixCacheTrie() + kv = _make_kv_list(3, seed=1) + trie.insert([1, 2, 3], kv) + + length, result = trie.lookup([1, 2, 3, 4, 5]) + assert length == 3 + assert result is kv + + def test_shorter_query_than_entry(self): + """Query shorter than any stored prefix should not match.""" + trie = PrefixCacheTrie() + trie.insert([1, 2, 3, 4, 5], _make_kv_list(5)) + + length, result = trie.lookup([1, 2]) + assert length == 0 + assert result is None + + def test_divergent_prefixes(self): + trie = PrefixCacheTrie() + kv_a = _make_kv_list(3, seed=10) + kv_b = _make_kv_list(3, seed=20) + trie.insert([1, 2, 3], kv_a) + trie.insert([1, 2, 4], kv_b) + + len_a, res_a = trie.lookup([1, 2, 3]) + assert len_a == 3 + assert res_a is kv_a + + len_b, res_b = trie.lookup([1, 2, 4]) + assert len_b == 3 + assert res_b is kv_b + + def test_multiple_prefixes_share_common_root(self): + trie = PrefixCacheTrie() + kv1 = _make_kv_list(2, seed=1) + kv2 = _make_kv_list(4, seed=2) + kv3 = _make_kv_list(6, seed=3) + trie.insert([10, 20], kv1) + trie.insert([10, 20, 30, 40], kv2) + trie.insert([10, 20, 30, 40, 50, 60], kv3) + + length, result = trie.lookup([10, 20, 30, 40, 50]) + assert length == 4 + assert result is kv2 + + +class TestLRUEviction: + def test_evict_reduces_entries(self): + trie = PrefixCacheTrie(max_entries=2) + trie.insert([1], _make_kv_list(1, seed=1)) + time.sleep(0.01) + trie.insert([2], _make_kv_list(1, seed=2)) + time.sleep(0.01) + trie.insert([3], _make_kv_list(1, seed=3)) + + stats = trie.get_stats() + assert stats["entries"] <= 2 + + def test_evict_removes_oldest(self): + trie = PrefixCacheTrie(max_entries=100) + kv_old = _make_kv_list(1, seed=1) + kv_new = _make_kv_list(1, seed=2) + trie.insert([1], kv_old) + time.sleep(0.01) + trie.insert([2], kv_new) + + trie.evict_lru(1) + + _, res_old = trie.lookup([1]) + _, res_new = trie.lookup([2]) + assert res_old is None + assert res_new is not None + + def test_auto_eviction_on_insert(self): + trie = PrefixCacheTrie(max_entries=3) + for i in range(5): + trie.insert([i * 10], _make_kv_list(1, seed=i)) + time.sleep(0.005) + + stats = trie.get_stats() + assert stats["entries"] <= 3 + + +class TestRadixCompression: + def test_single_child_chains_collapsed(self): + trie = PrefixCacheTrie() + trie.insert([1, 2, 3, 4, 5], _make_kv_list(5, seed=1)) + + # After insert the path should be a single compressed node + assert 1 in trie._root.children + child = trie._root.children[1] + assert child.tokens == [1, 2, 3, 4, 5] + + def test_split_preserves_data(self): + trie = PrefixCacheTrie() + kv_long = _make_kv_list(5, seed=1) + kv_short = _make_kv_list(3, seed=2) + + trie.insert([1, 2, 3, 4, 5], kv_long) + trie.insert([1, 2, 3], kv_short) + + len_short, res_short = trie.lookup([1, 2, 3]) + assert len_short == 3 + assert res_short is kv_short + + len_long, res_long = trie.lookup([1, 2, 3, 4, 5]) + assert len_long == 5 + assert res_long is kv_long + + def test_compression_after_eviction(self): + trie = PrefixCacheTrie(max_entries=100) + trie.insert([1, 2, 3], _make_kv_list(3, seed=1)) + time.sleep(0.01) + trie.insert([1, 2, 3, 4, 5], _make_kv_list(5, seed=2)) + + # Evict the shorter one, then the chain [1,2,3] -> [4,5] should compress + trie.evict_lru(1) + + len_long, res_long = trie.lookup([1, 2, 3, 4, 5]) + assert len_long == 5 + assert res_long is not None + + +class TestStats: + def test_initial_stats(self): + trie = PrefixCacheTrie() + stats = trie.get_stats() + assert stats["entries"] == 0 + assert stats["total_lookups"] == 0 + assert stats["total_hits"] == 0 + assert stats["hit_rate"] == 0.0 + + def test_hit_rate(self): + trie = PrefixCacheTrie() + trie.insert([1, 2, 3], _make_kv_list(3)) + + trie.lookup([1, 2, 3]) + trie.lookup([9, 9, 9]) + + stats = trie.get_stats() + assert stats["total_lookups"] == 2 + assert stats["total_hits"] == 1 + assert abs(stats["hit_rate"] - 0.5) < 1e-6 + + def test_memory_bytes(self): + trie = PrefixCacheTrie() + kv = _make_kv_list(10) + trie.insert([1, 2, 3], kv) + + stats = trie.get_stats() + assert stats["memory_bytes"] > 0 + + +class TestThreadSafety: + def test_concurrent_inserts(self): + trie = PrefixCacheTrie(max_entries=200) + errors: list[Exception] = [] + + def inserter(start): + try: + for i in range(20): + tokens = [start * 1000 + i] + trie.insert(tokens, _make_kv_list(1, seed=start * 100 + i)) + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=inserter, args=(t,)) for t in range(5)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors + stats = trie.get_stats() + assert stats["entries"] <= 200 + assert stats["entries"] > 0 + + def test_concurrent_lookups(self): + trie = PrefixCacheTrie() + trie.insert([1, 2, 3], _make_kv_list(3)) + errors: list[Exception] = [] + + def reader(): + try: + for _ in range(50): + length, _ = trie.lookup([1, 2, 3]) + assert length == 3 + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=reader) for _ in range(5)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors + + +class TestGetOrComputePrefix: + def test_cache_miss_calls_compute(self): + trie = PrefixCacheTrie() + called = [False] + + def compute(tokens): + called[0] = True + return _make_kv_list(len(tokens)) + + hit_len, kv = get_or_compute_prefix(trie, [1, 2, 3], compute, NUM_LAYERS) + assert hit_len == 0 + assert called[0] + assert kv is not None + + def test_cache_hit_skips_compute(self): + trie = PrefixCacheTrie() + kv_original = _make_kv_list(3, seed=42) + trie.insert([1, 2, 3], kv_original) + called = [False] + + def compute(tokens): + called[0] = True + return _make_kv_list(len(tokens)) + + hit_len, kv = get_or_compute_prefix(trie, [1, 2, 3], compute, NUM_LAYERS) + assert hit_len == 3 + assert not called[0] + assert kv is kv_original + + def test_stores_result_for_later(self): + trie = PrefixCacheTrie() + + def compute(tokens): + return _make_kv_list(len(tokens), seed=99) + + get_or_compute_prefix(trie, [5, 6, 7], compute, NUM_LAYERS) + + length, result = trie.lookup([5, 6, 7]) + assert length == 3 + assert result is not None diff --git a/tests/test_ssd_kv_cache.py b/tests/test_ssd_kv_cache.py new file mode 100644 index 0000000..9cd7d27 --- /dev/null +++ b/tests/test_ssd_kv_cache.py @@ -0,0 +1,313 @@ +"""Tests for SSD KV cache persistence (hot/cold 2-tier cache).""" + +from __future__ import annotations + +import tempfile +from pathlib import Path + +import pytest + +try: + import mlx.core as mx + + HAS_MLX = True +except ImportError: + HAS_MLX = False + +from mlx_flash_compress.kv_cache_backend import ( + KVCacheBackend, + PlainKVCache, + create_kv_cache, +) +from mlx_flash_compress.ssd_kv_cache import SSDKVCache + +pytestmark = pytest.mark.skipif(not HAS_MLX, reason="MLX required") + +NUM_LAYERS = 4 +NUM_HEADS = 2 +HEAD_DIM = 64 + + +def _make_kv(seq_len, num_heads=NUM_HEADS, head_dim=HEAD_DIM, seed=0): + mx.random.seed(seed) + keys = mx.random.normal((num_heads, seq_len, head_dim)) + values = mx.random.normal((num_heads, seq_len, head_dim)) + return keys, values + + +class TestSSDKVCacheBasic: + def test_is_kv_cache_backend(self): + with tempfile.TemporaryDirectory() as td: + cache = SSDKVCache(NUM_LAYERS, NUM_HEADS, HEAD_DIM, cache_dir=td) + assert isinstance(cache, KVCacheBackend) + + def test_update_and_get_within_ram(self): + with tempfile.TemporaryDirectory() as td: + cache = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + max_ram_tokens=100, + cache_dir=td, + ) + k, v = _make_kv(10) + all_k, all_v = cache.update(0, k, v) + assert all_k.shape == (NUM_HEADS, 10, HEAD_DIM) + assert all_v.shape == (NUM_HEADS, 10, HEAD_DIM) + + def test_get_kv_empty(self): + with tempfile.TemporaryDirectory() as td: + cache = SSDKVCache(NUM_LAYERS, NUM_HEADS, HEAD_DIM, cache_dir=td) + k, v = cache.get_kv(0) + assert k.shape == (NUM_HEADS, 0, HEAD_DIM) + + def test_stats_initial(self): + with tempfile.TemporaryDirectory() as td: + cache = SSDKVCache(NUM_LAYERS, NUM_HEADS, HEAD_DIM, cache_dir=td) + stats = cache.get_stats() + assert stats["strategy"] == "ssd" + assert stats["ssd_writes"] == 0 + assert stats["ssd_reads"] == 0 + + +class TestSSDEviction: + def test_eviction_to_ssd(self): + with tempfile.TemporaryDirectory() as td: + max_ram = 8 + cache = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + max_ram_tokens=max_ram, + cache_dir=td, + ) + k, v = _make_kv(12, seed=1) + cache.update(0, k, v) + + hot_k, hot_v = cache._hot.get_kv(0) + assert hot_k.shape[1] <= max_ram + + stats = cache.get_stats() + assert stats["ssd_writes"] >= 1 + assert stats["bytes_on_ssd"] > 0 + + def test_get_kv_merges_hot_and_cold(self): + with tempfile.TemporaryDirectory() as td: + max_ram = 5 + cache = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + max_ram_tokens=max_ram, + cache_dir=td, + ) + k, v = _make_kv(10, seed=2) + cache.update(0, k, v) + + all_k, all_v = cache.get_kv(0) + assert all_k.shape[1] == 10 + + def test_multiple_evictions_accumulate(self): + with tempfile.TemporaryDirectory() as td: + max_ram = 4 + cache = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + max_ram_tokens=max_ram, + cache_dir=td, + ) + for i in range(3): + k, v = _make_kv(5, seed=i) + cache.update(0, k, v) + + all_k, _ = cache.get_kv(0) + assert all_k.shape[1] == 15 + + stats = cache.get_stats() + assert stats["ssd_writes"] >= 2 + + +class TestSSDRoundTrip: + def test_write_read_data_integrity(self): + """Values written to SSD and read back must match the originals.""" + with tempfile.TemporaryDirectory() as td: + max_ram = 4 + cache = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + max_ram_tokens=max_ram, + cache_dir=td, + ) + k, v = _make_kv(8, seed=42) + cache.update(0, k, v) + + all_k, all_v = cache.get_kv(0) + assert mx.allclose(all_k, k, atol=1e-5) + assert mx.allclose(all_v, v, atol=1e-5) + + def test_safetensors_files_created(self): + with tempfile.TemporaryDirectory() as td: + cache = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + max_ram_tokens=4, + cache_dir=td, + ) + k, v = _make_kv(10, seed=3) + cache.update(0, k, v) + + sf_files = list(Path(td).rglob("*.safetensors")) + assert len(sf_files) >= 1 + + +class TestSSDSessionIsolation: + def test_different_prefixes_get_different_dirs(self): + with tempfile.TemporaryDirectory() as td: + c1 = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + cache_dir=td, + prompt_prefix="session_alpha", + ) + c2 = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + cache_dir=td, + prompt_prefix="session_beta", + ) + assert c1._session_dir != c2._session_dir + + def test_sessions_do_not_share_data(self): + with tempfile.TemporaryDirectory() as td: + c1 = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + max_ram_tokens=4, + cache_dir=td, + prompt_prefix="a", + ) + c2 = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + max_ram_tokens=4, + cache_dir=td, + prompt_prefix="b", + ) + k, v = _make_kv(8, seed=10) + c1.update(0, k, v) + + k2, v2 = c2.get_kv(0) + assert k2.shape[1] == 0 + + +class TestSSDReset: + def test_reset_clears_hot_and_cold(self): + with tempfile.TemporaryDirectory() as td: + cache = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + max_ram_tokens=4, + cache_dir=td, + ) + k, v = _make_kv(10, seed=5) + cache.update(0, k, v) + + cache.reset() + + rk, rv = cache.get_kv(0) + assert rk.shape[1] == 0 + + stats = cache.get_stats() + assert stats["ssd_writes"] == 0 + assert stats["bytes_on_ssd"] == 0 + + def test_reset_removes_ssd_files(self): + with tempfile.TemporaryDirectory() as td: + cache = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + max_ram_tokens=4, + cache_dir=td, + ) + k, v = _make_kv(10, seed=6) + cache.update(0, k, v) + session_dir = cache._session_dir + + cache.reset() + + sf_files = list(session_dir.rglob("*.safetensors")) + assert len(sf_files) == 0 + + +class TestSSDMultiLayer: + def test_independent_layer_eviction(self): + with tempfile.TemporaryDirectory() as td: + cache = SSDKVCache( + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + max_ram_tokens=4, + cache_dir=td, + ) + for layer in range(NUM_LAYERS): + k, v = _make_kv(6 + layer, seed=layer) + cache.update(layer, k, v) + + for layer in range(NUM_LAYERS): + all_k, _ = cache.get_kv(layer) + assert all_k.shape[1] == 6 + layer + + +class TestSSDFactory: + def test_create_ssd(self): + with tempfile.TemporaryDirectory() as td: + backend = create_kv_cache( + "ssd", + num_layers=4, + num_heads=2, + head_dim=64, + max_ram_tokens=16, + cache_dir=td, + ) + assert isinstance(backend, SSDKVCache) + + def test_create_ssd_hybrid(self): + with tempfile.TemporaryDirectory() as td: + backend = create_kv_cache( + "ssd_hybrid", + num_layers=4, + num_heads=2, + head_dim=64, + max_ram_tokens=16, + cache_dir=td, + key_bits=4, + value_bits=4, + calibration_tokens=0, + ) + assert isinstance(backend, SSDKVCache) + + def test_ssd_hybrid_inner_is_quantized(self): + from mlx_flash_compress.kv_cache_backend import QuantizedKVCache + + with tempfile.TemporaryDirectory() as td: + backend = create_kv_cache( + "ssd_hybrid", + num_layers=4, + num_heads=2, + head_dim=64, + max_ram_tokens=16, + cache_dir=td, + key_bits=4, + value_bits=4, + calibration_tokens=0, + ) + assert isinstance(backend._hot, QuantizedKVCache) diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py new file mode 100644 index 0000000..abd10cf --- /dev/null +++ b/tests/test_telemetry.py @@ -0,0 +1,330 @@ +"""Tests for hardware telemetry module.""" + +from __future__ import annotations + +import time +from unittest.mock import MagicMock, patch + +import pytest + +from mlx_flash_compress.telemetry import ( + _MAX_HISTORY, + HardwareTelemetry, + TelemetrySample, +) + + +class TestTelemetrySample: + def test_defaults(self): + s = TelemetrySample() + assert s.timestamp == 0.0 + assert s.gpu_util_pct == 0.0 + assert s.gpu_renderer_pct == 0.0 + assert s.gpu_tiler_pct == 0.0 + assert s.gpu_memory_used_gb == 0.0 + assert s.gpu_memory_total_gb == 0.0 + assert s.power_watts == 0.0 + assert s.cpu_temp_c == 0.0 + assert s.gpu_temp_c == 0.0 + assert s.memory_pressure == "normal" + assert s.memory_used_gb == 0.0 + assert s.memory_total_gb == 0.0 + assert s.ane_util_pct == 0.0 + + def test_custom_values(self): + s = TelemetrySample( + timestamp=1000.0, + gpu_util_pct=75.0, + power_watts=12.5, + cpu_temp_c=65.0, + gpu_temp_c=72.0, + memory_pressure="warn", + ) + assert s.gpu_util_pct == 75.0 + assert s.power_watts == 12.5 + assert s.cpu_temp_c == 65.0 + assert s.memory_pressure == "warn" + + def test_to_dict(self): + s = TelemetrySample(timestamp=1.0, gpu_util_pct=50.0) + d = s.to_dict() + assert isinstance(d, dict) + assert d["timestamp"] == 1.0 + assert d["gpu_util_pct"] == 50.0 + assert d["memory_pressure"] == "normal" + assert len(d) == 13 # all fields + + +# Reusable mock ioreg output for GPU metrics +MOCK_IOREG_GPU = """ ++-o IOAccelerator + | "Device Utilization %" = 42 + | "Renderer Utilization %" = 38 + | "Tiler Utilization %" = 15 + | "In use system memory"= 2147483648 + | "Alloc system memory"= 4294967296 +""" + +MOCK_IOREG_BATTERY = """ ++-o AppleSmartBattery + | "InstantAmperage"= 1500 + | "Voltage"= 12800 +""" + +MOCK_VM_STAT = """Mach Virtual Memory Statistics: (page size of 16384 bytes) +Pages free: 50000. +Pages active: 200000. +Pages inactive: 100000. +Pages speculative: 5000. +Pages throttled: 0. +Pages wired down: 80000. +Pages purgeable: 10000. +"Translation faults": 500000000. +Pages copy-on-write: 30000000. +Pages zero filled: 200000000. +Pages reactivated: 1000000. +Pages purged: 2000000. +File-backed pages: 120000. +Anonymous pages: 80000. +Pages stored in compressor: 30000. +Pages occupied by compressor: 15000. +""" + + +class TestHardwareTelemetry: + def _make_telemetry(self): + """Create a HardwareTelemetry with mocked total memory detection.""" + with patch.object(HardwareTelemetry, "_detect_total_memory", return_value=32.0): + return HardwareTelemetry() + + @patch("mlx_flash_compress.telemetry.HardwareTelemetry._run_cmd") + def test_sample_with_gpu_metrics(self, mock_cmd): + """Test that GPU metrics are parsed from ioreg output.""" + + def side_effect(args, timeout=3.0): + if "IOAccelerator" in args: + return MOCK_IOREG_GPU + return "" + + mock_cmd.side_effect = side_effect + tel = self._make_telemetry() + s = tel.sample() + + assert s.gpu_util_pct == 42.0 + assert s.gpu_renderer_pct == 38.0 + assert s.gpu_tiler_pct == 15.0 + assert s.gpu_memory_used_gb == pytest.approx(2.0, abs=0.1) + assert s.gpu_memory_total_gb == pytest.approx(4.0, abs=0.1) + + @patch("mlx_flash_compress.telemetry.HardwareTelemetry._run_cmd") + def test_sample_with_battery_power(self, mock_cmd): + """Test power extraction from battery info.""" + + def side_effect(args, timeout=3.0): + if "AppleSmartBattery" in args: + return MOCK_IOREG_BATTERY + return "" + + mock_cmd.side_effect = side_effect + tel = self._make_telemetry() + s = tel.sample() + + # 1500 mA * 12800 mV = 19,200,000 / 1,000,000 = 19.2 W + assert s.power_watts == pytest.approx(19.2, abs=0.1) + + @patch("mlx_flash_compress.telemetry.HardwareTelemetry._run_cmd") + def test_sample_with_vm_stat(self, mock_cmd): + """Test memory parsing from vm_stat output.""" + + def side_effect(args, timeout=3.0): + if args == ["vm_stat"]: + return MOCK_VM_STAT + if args == ["memory_pressure"]: + return "System-wide memory free percentage: 50%" + return "" + + mock_cmd.side_effect = side_effect + + # Also mock psutil as not available + with patch.dict("sys.modules", {"psutil": None}): + tel = self._make_telemetry() + s = tel.sample() + + assert s.memory_total_gb == 32.0 + # used = (active + wired + compressed) * page_size / 1GB + # = (200000 + 80000 + 15000) * 16384 / 1073741824 + expected_used = (200000 + 80000 + 15000) * 16384 / (1024**3) + assert s.memory_used_gb == pytest.approx(expected_used, abs=0.5) + + @patch("mlx_flash_compress.telemetry.HardwareTelemetry._run_cmd") + def test_graceful_failure_all_commands(self, mock_cmd): + """Test that sample returns zeros when all commands fail.""" + mock_cmd.return_value = "" + tel = self._make_telemetry() + s = tel.sample() + + assert s.gpu_util_pct == 0.0 + assert s.power_watts == 0.0 + assert s.cpu_temp_c == 0.0 + assert s.gpu_temp_c == 0.0 + assert s.memory_pressure == "normal" + assert s.timestamp > 0 # timestamp always set + + @patch("mlx_flash_compress.telemetry.HardwareTelemetry._run_cmd") + def test_circular_buffer_history(self, mock_cmd): + """Test that history is bounded by _MAX_HISTORY.""" + mock_cmd.return_value = "" + tel = self._make_telemetry() + + # Insert more samples than buffer size + for i in range(_MAX_HISTORY + 20): + tel.sample() + + history = tel.get_history(seconds=9999) + assert len(history) == _MAX_HISTORY + + @patch("mlx_flash_compress.telemetry.HardwareTelemetry._run_cmd") + def test_get_history_time_filter(self, mock_cmd): + """Test that get_history respects the seconds parameter.""" + mock_cmd.return_value = "" + tel = self._make_telemetry() + + # Manually insert samples with old timestamps + old_sample = TelemetrySample(timestamp=time.time() - 300) # 5 min ago + recent_sample = TelemetrySample(timestamp=time.time() - 10) # 10s ago + tel._history.append(old_sample) + tel._history.append(recent_sample) + + history = tel.get_history(seconds=60) + assert len(history) == 1 # only the recent one + + @patch("mlx_flash_compress.telemetry.HardwareTelemetry._run_cmd") + def test_get_stats_empty(self, mock_cmd): + """Test get_stats with no samples.""" + mock_cmd.return_value = "" + tel = self._make_telemetry() + stats = tel.get_stats() + + assert stats["samples_count"] == 0 + assert "current" in stats + assert stats["current"]["gpu_util_pct"] == 0.0 + + @patch("mlx_flash_compress.telemetry.HardwareTelemetry._run_cmd") + def test_get_stats_aggregation(self, mock_cmd): + """Test that get_stats computes correct min/max/avg.""" + mock_cmd.return_value = "" + tel = self._make_telemetry() + + # Insert samples with known values + for gpu_pct in [10.0, 20.0, 30.0, 40.0, 50.0]: + s = TelemetrySample(timestamp=time.time(), gpu_util_pct=gpu_pct) + tel._history.append(s) + + stats = tel.get_stats() + assert stats["samples_count"] == 5 + assert stats["avg"]["gpu_util_pct"] == 30.0 + assert stats["max"]["gpu_util_pct"] == 50.0 + assert stats["min"]["gpu_util_pct"] == 10.0 + + @patch("mlx_flash_compress.telemetry.HardwareTelemetry._run_cmd") + def test_start_stop_sampling(self, mock_cmd): + """Test background sampling thread lifecycle.""" + mock_cmd.return_value = "" + tel = self._make_telemetry() + + assert not tel._sampling + tel.start_sampling(interval_ms=50) + assert tel._sampling + assert tel._thread is not None + assert tel._thread.is_alive() + + # Let it run a few cycles + time.sleep(0.2) + tel.stop_sampling() + + assert not tel._sampling + # Should have collected some samples + assert len(tel._history) > 0 + + @patch("mlx_flash_compress.telemetry.HardwareTelemetry._run_cmd") + def test_start_sampling_idempotent(self, mock_cmd): + """Test that calling start_sampling twice doesn't create two threads.""" + mock_cmd.return_value = "" + tel = self._make_telemetry() + + tel.start_sampling(interval_ms=100) + thread1 = tel._thread + tel.start_sampling(interval_ms=100) + thread2 = tel._thread + + assert thread1 is thread2 + tel.stop_sampling() + + @patch("mlx_flash_compress.telemetry.HardwareTelemetry._run_cmd") + def test_sample_thread_safety(self, mock_cmd): + """Test concurrent access to sample history.""" + mock_cmd.return_value = "" + tel = self._make_telemetry() + + import threading + + errors = [] + + def reader(): + try: + for _ in range(50): + tel.get_history(seconds=120) + tel.get_stats() + except Exception as e: + errors.append(e) + + def writer(): + try: + for _ in range(50): + tel.sample() + except Exception as e: + errors.append(e) + + threads = [ + threading.Thread(target=reader), + threading.Thread(target=reader), + threading.Thread(target=writer), + threading.Thread(target=writer), + ] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(errors) == 0 + + def test_extract_ioreg_int(self): + """Test the ioreg integer extraction helper.""" + tel = self._make_telemetry() + + text = ' "Device Utilization %" = 42\n "Renderer Utilization %" = 38' + assert tel._extract_ioreg_int(text, "Device Utilization %") == 42 + assert tel._extract_ioreg_int(text, "Renderer Utilization %") == 38 + assert tel._extract_ioreg_int(text, "NonExistent") == 0 + + def test_extract_ioreg_int_large_value(self): + """Test extraction of large byte values.""" + tel = self._make_telemetry() + text = ' "In use system memory"= 2147483648' + assert tel._extract_ioreg_int(text, '"In use system memory"=') == 2147483648 + + def test_detect_total_memory(self): + """Test total memory detection via sysctl.""" + with patch.object( + HardwareTelemetry, + "_run_cmd", + return_value="34359738368\n", + ): + tel = HardwareTelemetry() + assert tel._total_memory_gb == pytest.approx(32.0, abs=0.1) + + def test_detect_total_memory_failure(self): + """Test graceful failure when sysctl is unavailable.""" + with patch.object(HardwareTelemetry, "_run_cmd", return_value=""): + tel = HardwareTelemetry() + assert tel._total_memory_gb == 0.0 From d1be4066754d0589c02bf32a25ec9d9163119d5a Mon Sep 17 00:00:00 2001 From: szibis <329831+szibis@users.noreply.github.com> Date: Mon, 11 May 2026 14:34:37 +0200 Subject: [PATCH 2/6] =?UTF-8?q?feat:=20overhaul=20chat=20UI=20and=20dashbo?= =?UTF-8?q?ard=20=E2=80=94=20dynamic=20models,=20SSE=20streaming,=20teleme?= =?UTF-8?q?try=20integration?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Chat UI: fetch models from /v1/models/registry instead of hardcoded list, SSE streaming with token-by-token rendering, optgroup categories, 502/503 error states. Dashboard: 10-link nav bar, API endpoints card, inline power/temp/ANE in GPU card, worker session details. --- mlx-flash-server/src/chat_ui.rs | 234 ++++++++++++++++++++++++------ mlx-flash-server/src/dashboard.rs | 96 +++++++++++- 2 files changed, 278 insertions(+), 52 deletions(-) diff --git a/mlx-flash-server/src/chat_ui.rs b/mlx-flash-server/src/chat_ui.rs index 1915eee..d405e38 100644 --- a/mlx-flash-server/src/chat_ui.rs +++ b/mlx-flash-server/src/chat_ui.rs @@ -29,12 +29,18 @@ const CHAT_HTML: &str = r##" .topbar { padding: 14px 24px; border-bottom: 1px solid var(--border); display: flex; align-items: center; gap: 12px; flex-shrink: 0; } .topbar h1 { font-size: 1.1rem; font-weight: 600; background: linear-gradient(135deg, var(--accent), var(--accent2)); -webkit-background-clip: text; -webkit-text-fill-color: transparent; } .topbar .model-name { font-size: 0.8rem; color: var(--dim); padding: 3px 10px; background: var(--card); border-radius: 6px; border: 1px solid var(--border); } - .topbar .nav { margin-left: auto; display: flex; gap: 8px; } + .topbar .nav { margin-left: auto; display: flex; gap: 8px; flex-wrap: wrap; } .topbar .nav a { color: var(--dim); text-decoration: none; font-size: 0.8rem; padding: 4px 10px; border-radius: 6px; transition: all 0.15s; } .topbar .nav a:hover { color: var(--text); background: var(--hover); } + .topbar .nav a.nav-active { background: rgba(77,166,255,0.15); color: var(--accent); } .model-select { background: var(--card); color: var(--text); border: 1px solid var(--border); border-radius: 8px; padding: 5px 12px; font-size: 0.8rem; font-family: inherit; cursor: pointer; outline: none; -webkit-appearance: none; max-width: 280px; } .model-select:focus { border-color: var(--accent); } .model-select option { background: var(--card); color: var(--text); } + .model-select optgroup { background: var(--surface); color: var(--dim); font-style: normal; font-size: 0.75rem; } + .loading-indicator { display: none; align-items: center; gap: 8px; padding: 12px 0; } + .loading-indicator.active { display: flex; } + .loading-spinner { width: 16px; height: 16px; border: 2px solid var(--border); border-top-color: var(--accent); border-radius: 50%; animation: spin 0.8s linear infinite; } + @keyframes spin { to { transform: rotate(360deg); } } .messages { flex: 1; overflow-y: auto; padding: 20px 0; scroll-behavior: smooth; } .messages::-webkit-scrollbar { width: 6px; } @@ -120,7 +126,11 @@ const CHAT_HTML: &str = r##"
Dashboard - Chat + Chat + Prometheus + Status + GPU + Workers
@@ -476,11 +486,11 @@ async function sendMessage() { try { const sel = document.getElementById('model-select'); - const currentModel = sel.value || localStorage.getItem('mlx-flash-model') || 'local'; + const currentModel = sel.value || activeModelId || localStorage.getItem('mlx-flash-model') || 'local'; const payload = { model: currentModel, messages: messages.map(m => ({role: m.role, content: m.content})), - stream: false, max_tokens: 1024, + stream: true, max_tokens: 1024, }; const t0 = Date.now(); @@ -489,12 +499,9 @@ async function sendMessage() { headers: {'Content-Type': 'application/json'}, body: JSON.stringify(payload), }); - const elapsed = ((Date.now() - t0) / 1000).toFixed(1); - - let fullText = ''; - let respTokens = 0; if (!resp.ok) { + const elapsed = ((Date.now() - t0) / 1000).toFixed(1); let errMsg = 'Server error ' + resp.status; try { const j = await resp.json(); errMsg = j.error || errMsg; } catch(e) {} @@ -503,6 +510,10 @@ async function sendMessage() { + '

Start it in another terminal:

' + '
'
           + 'pip install mlx mlx-lm mlx-flash\nmlx-flash --port 8081 --preload
'; + } else if (resp.status === 503 || errMsg.toLowerCase().includes('loading')) { + aiEl.innerHTML = '

Model is still loading...

' + + '

Please wait for the model to finish loading, then try again. ' + + 'Check the Dashboard for progress.

'; } else { aiEl.innerHTML = '

' + errMsg + '

'; } @@ -510,25 +521,76 @@ async function sendMessage() { return; } + // SSE streaming — read tokens as they arrive + let fullText = ''; + let respTokens = 0; let tokPerSec = ''; - try { - const json = await resp.json(); - fullText = json.choices?.[0]?.message?.content - || json.choices?.[0]?.text - || JSON.stringify(json); - // Extract MLX-Flash performance data if present - const mlxData = json.mlx_flash_compress || {}; - if (mlxData.tok_per_s) tokPerSec = mlxData.tok_per_s.toFixed(1) + ' tok/s'; - const usage = json.usage || {}; - respTokens = usage.completion_tokens || 0; - if (respTokens && !tokPerSec) { - tokPerSec = (respTokens / parseFloat(elapsed)).toFixed(1) + ' tok/s'; + let lastChunkModel = currentModel; + + const contentType = resp.headers.get('content-type') || ''; + if (contentType.includes('text/event-stream') || contentType.includes('text/plain')) { + // SSE stream mode + const reader = resp.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; + + while (true) { + const {done, value} = await reader.read(); + if (done) break; + buffer += decoder.decode(value, {stream: true}); + const lines = buffer.split('\n'); + buffer = lines.pop() || ''; + + for (const line of lines) { + if (!line.startsWith('data: ')) continue; + const payload = line.slice(6).trim(); + if (payload === '[DONE]') continue; + try { + const chunk = JSON.parse(payload); + const delta = chunk.choices?.[0]?.delta || {}; + if (delta.content) { + fullText += delta.content; + aiEl.innerHTML = formatContent(fullText); + messagesEl.scrollTop = messagesEl.scrollHeight; + } + if (chunk.model) lastChunkModel = chunk.model; + // Some implementations send usage in the final chunk + if (chunk.usage) { + respTokens = chunk.usage.completion_tokens || 0; + } + // Extract MLX-Flash performance data if present + const mlxData = chunk.mlx_flash_compress || {}; + if (mlxData.tok_per_s) tokPerSec = mlxData.tok_per_s.toFixed(1) + ' tok/s'; + } catch(e) { /* skip unparseable lines */ } + } } - } catch(e) { - fullText = 'Failed to parse response'; + } else { + // Fallback: non-streaming JSON response + try { + const json = await resp.json(); + fullText = json.choices?.[0]?.message?.content + || json.choices?.[0]?.text + || JSON.stringify(json); + const mlxData = json.mlx_flash_compress || {}; + if (mlxData.tok_per_s) tokPerSec = mlxData.tok_per_s.toFixed(1) + ' tok/s'; + const usage = json.usage || {}; + respTokens = usage.completion_tokens || 0; + if (json.model) lastChunkModel = json.model; + } catch(e) { + fullText = 'Failed to parse response'; + } + aiEl.innerHTML = formatContent(fullText); + } + + const elapsed = ((Date.now() - t0) / 1000).toFixed(1); + // Estimate tokens from text if not provided + if (!respTokens && fullText) { + respTokens = Math.ceil(fullText.length / 4); + } + if (respTokens && !tokPerSec) { + tokPerSec = (respTokens / parseFloat(elapsed)).toFixed(1) + ' tok/s'; } - aiEl.innerHTML = formatContent(fullText); messages[messages.length-1].content = fullText; // Track cumulative tokens for savings calculator if (respTokens > 0) { @@ -539,11 +601,11 @@ async function sendMessage() { // Get current model + pressure for metadata bar let currentPressure = 'Normal'; - let currentModelName = currentModel; + let currentModelName = lastChunkModel || currentModel; try { const st = await fetch('/status').then(r=>r.json()); currentPressure = (st.memory||{}).pressure || 'Normal'; - currentModelName = st.model || currentModel; + if (st.model) currentModelName = st.model; } catch(e) {} // Set metadata bar under the AI message @@ -570,31 +632,97 @@ async function sendMessage() { } } -// Known models catalog (matches Python MODELS list) -const KNOWN_MODELS = [ - {id:'mlx-community/gemma-4-E2B-it-4bit', label:'Gemma 4 E2B (1.5GB)', size:1.5}, - {id:'mlx-community/gemma-4-E4B-it-4bit', label:'Gemma 4 E4B (2.8GB)', size:2.8}, - {id:'mlx-community/Qwen3-4B-4bit', label:'Qwen3 4B (2.5GB)', size:2.5}, - {id:'mlx-community/Qwen3-8B-4bit', label:'Qwen3 8B (5GB)', size:5.0}, - {id:'mlx-community/gemma-4-26b-it-4bit', label:'Gemma 4 26B MoE (15GB)', size:15.0}, - {id:'mlx-community/Qwen3-30B-A3B-4bit', label:'Qwen3 30B MoE (18GB)', size:18.0}, - {id:'mlx-community/gemma-4-31b-it-4bit', label:'Gemma 4 31B (20GB)', size:20.0}, - {id:'mlx-community/Mixtral-8x7B-Instruct-v0.1-4bit', label:'Mixtral 8x7B (26GB)', size:26.0}, -]; +// Dynamic model registry — populated from /v1/models/registry at page load +let KNOWN_MODELS = []; +let activeModelId = ''; + +async function loadModelRegistry() { + try { + const resp = await fetch('/v1/models/registry'); + const data = await resp.json(); + const models = data.models || data || []; + KNOWN_MODELS = models.map(m => ({ + id: m.id || m.model_id || m.name, + label: (m.id || m.name || '').split('/').pop() + (m.size_gb_approx ? ' (' + m.size_gb_approx + 'GB)' : ''), + size: m.size_gb_approx || 0, + category: m.category_expected || '', + })); + } catch(e) { + // Fallback to a minimal list if registry fails + KNOWN_MODELS = [ + {id:'mlx-community/Llama-3.2-3B-Instruct-4bit', label:'Llama 3.2 3B (2GB)', size:2, category:'small_dense'}, + {id:'mlx-community/Qwen3-30B-A3B-4bit', label:'Qwen3 30B MoE (18GB)', size:18, category:'large_moe'}, + {id:'mlx-community/Qwen3.6-35B-A3B-4bit', label:'Qwen3.6 35B MoE (20GB)', size:20, category:'large_moe'}, + ]; + } + // Also add the currently active model if not in list + try { + const st = await fetch('/status'); + const data = await st.json(); + if (data.model) { + activeModelId = data.model; + if (!KNOWN_MODELS.some(m => m.id === activeModelId)) { + KNOWN_MODELS.unshift({id: activeModelId, label: activeModelId.split('/').pop() + ' (active)', size: 0, category: ''}); + } + } + } catch(e) {} + populateModelSelect(activeModelId || localStorage.getItem('mlx-flash-model') || ''); +} + +const CATEGORY_LABELS = { + 'small_dense': 'Small Dense', + 'medium_dense': 'Medium Dense', + 'large_dense': 'Large Dense', + 'small_moe': 'Small MoE', + 'medium_moe': 'Medium MoE', + 'large_moe': 'Large MoE', + '': 'Other', +}; function populateModelSelect(currentModel) { const sel = document.getElementById('model-select'); + // Group models by category + const groups = {}; const isKnown = KNOWN_MODELS.some(m => m.id === currentModel); - let opts = ''; + // If the active model isn't in the registry, add it at the top + let topOpts = ''; if (currentModel && !isKnown) { const short = currentModel.split('/').pop() || currentModel; - opts += ''; + topOpts += ''; } - opts += KNOWN_MODELS.map(m => - '' - ).join(''); - opts += ''; - sel.innerHTML = opts; + KNOWN_MODELS.forEach(m => { + const cat = m.category || ''; + if (!groups[cat]) groups[cat] = []; + groups[cat].push(m); + }); + let html = topOpts; + // Sort categories: small_dense, medium_dense, large_dense, small_moe, medium_moe, large_moe, then other + const catOrder = ['small_dense','medium_dense','large_dense','small_moe','medium_moe','large_moe','']; + const sortedCats = Object.keys(groups).sort((a,b) => { + const ai = catOrder.indexOf(a), bi = catOrder.indexOf(b); + return (ai===-1?99:ai) - (bi===-1?99:bi); + }); + if (sortedCats.length > 1 || (sortedCats.length === 1 && sortedCats[0] !== '')) { + sortedCats.forEach(cat => { + const label = CATEGORY_LABELS[cat] || cat || 'Other'; + html += ''; + groups[cat].forEach(m => { + const isActive = m.id === currentModel; + const displayLabel = isActive ? m.label + ' *' : m.label; + html += ''; + }); + html += ''; + }); + } else { + // Flat list if all same category + KNOWN_MODELS.forEach(m => { + const isActive = m.id === currentModel; + const displayLabel = isActive ? m.label + ' *' : m.label; + html += ''; + }); + } + html += ''; + sel.innerHTML = html; } async function switchModel(modelId) { @@ -628,15 +756,31 @@ async function switchModel(modelId) { } // Poll model + memory for status bar + pressure detection +let registryLoaded = false; async function updateStatus() { try { const st = await fetch('/status').then(r=>r.json()); const model = st.model || ''; - populateModelSelect(model); + activeModelId = model; + // On subsequent polls, just update the select highlight (don't re-fetch registry) + if (registryLoaded) { + populateModelSelect(model); + } const mem = st.memory || {}; const avail = Math.max((mem.free_gb||0)+(mem.inactive_gb||0)*0.5, 0); const pressure = (mem.pressure||'Normal').toString(); document.getElementById('mem-info').textContent = avail.toFixed(1)+'GB free / '+((mem.total_gb||0).toFixed(0))+'GB'; + // Detect model loading state + if (!model || model === 'loading...' || model === 'none') { + document.getElementById('status-text').textContent = 'Waiting for model to load...'; + document.getElementById('status-dot').style.background = 'var(--yellow)'; + } else if (!generating) { + const dotEl = document.getElementById('status-dot'); + if (dotEl.style.background !== 'var(--green)') { + document.getElementById('status-text').textContent = 'Ready'; + dotEl.style.background = 'var(--green)'; + } + } // Cumulative savings if (totalTokensSession > 0) { const totalSaved = calcSavings(totalTokensSession); @@ -646,6 +790,8 @@ async function updateStatus() { if (!generating) updatePressureBanner(pressure, avail.toFixed(1), mem.swap_used_gb||0); } catch(e) {} } +// Load model registry once at startup, then poll status +loadModelRegistry().then(() => { registryLoaded = true; }); updateStatus(); setInterval(updateStatus, 5000); diff --git a/mlx-flash-server/src/dashboard.rs b/mlx-flash-server/src/dashboard.rs index 97a8a0b..10a2128 100644 --- a/mlx-flash-server/src/dashboard.rs +++ b/mlx-flash-server/src/dashboard.rs @@ -26,9 +26,10 @@ const DASHBOARD_HTML: &str = r##" .header .status { display: flex; align-items: center; gap: 6px; font-size: 0.8rem; color: var(--dim); } .header .dot { width: 6px; height: 6px; border-radius: 50%; background: var(--green); animation: pulse 2s infinite; } @keyframes pulse { 0%,100% { opacity: 1; } 50% { opacity: 0.4; } } - .header .nav { display: flex; gap: 8px; } + .header .nav { display: flex; gap: 8px; flex-wrap: wrap; } .header .nav a { color: var(--dim); text-decoration: none; font-size: 0.8rem; padding: 4px 10px; border-radius: 6px; transition: all 0.15s; } .header .nav a:hover { color: var(--text); background: rgba(77,166,255,0.1); } + .header .nav a.nav-active { background: rgba(77,166,255,0.15); color: var(--accent); } .header .uptime { margin-left: auto; font-size: 0.8rem; color: var(--dim); font-variant-numeric: tabular-nums; } .container { padding: 20px 32px; } @@ -97,7 +98,18 @@ const DASHBOARD_HTML: &str = r##"

MLX-Flash

Live
-
DashboardChatMetrics
+
+ Dashboard + Chat + Prometheus + Status JSON + GPU + Telemetry + Cache + Workers + Models + Logs JSON +
0:00
@@ -169,7 +181,10 @@ const DASHBOARD_HTML: &str = r##"
Workers
-
Sessions: 0 | Strategy: least-connections + cache-affinity
+
+ Sessions: 0 | Strategy: least-connections + cache-affinity +
+
@@ -179,12 +194,17 @@ const DASHBOARD_HTML: &str = r##"
-
GPU (Metal)
+
GPU (Metal) + Telemetry
--%
-- GB used
--% renderer / --% tiler
-
Power: -- W | ANE: --%
+
+ Power: -- W | ANE: --% +
+
+ CPU Temp: --°C | GPU Temp: --°C +
Optimization Hints
@@ -223,6 +243,43 @@ const DASHBOARD_HTML: &str = r##"
+ +
+
API Endpoints
+
+
+ Core
+ /status + /health + /hints +
+
+ Chat
+ /v1/chat/completions + /v1/models +
+
+ Cache
+ /cache/stats +
+
+ Workers
+ /workers +
+
+ Telemetry
+ /telemetry + /gpu +
+
+ Admin
+ /metrics + /logs/recent + /v1/config +
+
+
+
@@ -342,11 +399,27 @@ function renderWorkers(wdata) { + ''; }).join(''); document.getElementById('worker-summary').textContent = wdata.healthy_count + '/' + wdata.total_count + ' healthy'; - // Show sessions + // Show sessions and strategy + if (wdata.strategy) { + document.getElementById('worker-strategy').textContent = wdata.strategy; + } if (wdata.sessions && wdata.sessions.length > 0) { - document.getElementById('sessions').textContent = wdata.sessions_active + ' (' + wdata.sessions.map(s => s.session_id.substring(0,8)+'→:'+s.worker_port).join(', ') + ')'; + document.getElementById('sessions').textContent = wdata.sessions_active || wdata.sessions.length; + // Build per-worker session detail + const perWorker = {}; + wdata.sessions.forEach(s => { + const port = s.worker_port || 'unknown'; + if (!perWorker[port]) perWorker[port] = []; + perWorker[port].push(s.session_id ? s.session_id.substring(0,8) : '?'); + }); + const detailParts = Object.entries(perWorker).map(([port, ids]) => + ':' + port + ' [' + ids.join(', ') + ']' + ); + document.getElementById('worker-sessions-detail').innerHTML = + 'Session pins: ' + detailParts.join(' | ') + ''; } else { document.getElementById('sessions').textContent = wdata.sessions_active || '0'; + document.getElementById('worker-sessions-detail').innerHTML = ''; } } @@ -526,7 +599,14 @@ async function poll() { curAne = telem.ane_util_pct || 0; document.getElementById('gpu-power').textContent = curPower.toFixed(1); document.getElementById('gpu-ane').textContent = curAne.toFixed(0) + '%'; - // Color-code temperature display + // Color-code temperature inline in the GPU card + const cpuTempEl = document.getElementById('gpu-cpu-temp'); + const gpuTempEl = document.getElementById('gpu-gpu-temp'); + cpuTempEl.textContent = curCpuTemp.toFixed(0); + gpuTempEl.textContent = curGpuTemp.toFixed(0); + cpuTempEl.className = tempColor(curCpuTemp); + gpuTempEl.className = tempColor(curGpuTemp); + // Color-code temperature display in chart const maxTemp = Math.max(curCpuTemp, curGpuTemp); const tc = tempColor(maxTemp); document.getElementById('temp-chart-val').innerHTML = 'CPU '+curCpuTemp.toFixed(0)+'° / GPU '+curGpuTemp.toFixed(0)+'°C'; From b8d41feb2bebcf6a7e3fe98fdbb908059eaa3007 Mon Sep 17 00:00:00 2001 From: szibis <329831+szibis@users.noreply.github.com> Date: Mon, 11 May 2026 14:50:16 +0200 Subject: [PATCH 3/6] fix: add SO_REUSEADDR to prevent 'Address already in use' on restart --- mlx_flash_compress/serve.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mlx_flash_compress/serve.py b/mlx_flash_compress/serve.py index b21ec28..58bcf0f 100644 --- a/mlx_flash_compress/serve.py +++ b/mlx_flash_compress/serve.py @@ -60,6 +60,7 @@ class ThreadedHTTPServer(ThreadingMixIn, HTTPServer): """HTTP server that handles each request in a new thread.""" daemon_threads = True + allow_reuse_address = True class InferenceState: From 2f7c893006956a9ce3c7e220246f5aeab4f92265 Mon Sep 17 00:00:00 2001 From: szibis <329831+szibis@users.noreply.github.com> Date: Mon, 11 May 2026 15:00:30 +0200 Subject: [PATCH 4/6] =?UTF-8?q?feat:=20smart=20model=20registry=20?= =?UTF-8?q?=E2=80=94=20RAM=20detection,=20download=20progress,=20filter=20?= =?UTF-8?q?unsupported?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Registry detects system RAM (sysctl HW_MEMSIZE) and filters dense models that won't fit; MoE models always shown (expert streaming) - Removed unsupported architectures (DeepSeek V4, gemma-4-e2b) from registry - Added 8 new supported models (Llama 1B/70B, Qwen3 4B/8B, Phi-4, gemma-4-12b, Mixtral) - Model switch shows real-time progress (downloading/loading/compiling phases) via /switch/progress polling endpoint - Chat UI: cached models marked with checkmark, MoE tag, category size hints - Python server: SO_REUSEADDR, switch progress tracking, phase updates during load --- mlx-flash-server/src/chat_ui.rs | 121 ++++++++++++++++++++++++-------- mlx-flash-server/src/server.rs | 56 ++++++++++++++- mlx_flash_compress/serve.py | 42 +++++++++++ scripts/models.yaml | 74 +++++++++++-------- 4 files changed, 234 insertions(+), 59 deletions(-) diff --git a/mlx-flash-server/src/chat_ui.rs b/mlx-flash-server/src/chat_ui.rs index d405e38..c937e3a 100644 --- a/mlx-flash-server/src/chat_ui.rs +++ b/mlx-flash-server/src/chat_ui.rs @@ -635,56 +635,67 @@ async function sendMessage() { // Dynamic model registry — populated from /v1/models/registry at page load let KNOWN_MODELS = []; let activeModelId = ''; +let systemRamGb = 0; async function loadModelRegistry() { try { const resp = await fetch('/v1/models/registry'); const data = await resp.json(); + systemRamGb = data.system_ram_gb || 0; const models = data.models || data || []; - KNOWN_MODELS = models.map(m => ({ + KNOWN_MODELS = models.filter(m => m.fits_ram !== false).map(m => ({ id: m.id || m.model_id || m.name, - label: (m.id || m.name || '').split('/').pop() + (m.size_gb_approx ? ' (' + m.size_gb_approx + 'GB)' : ''), - size: m.size_gb_approx || 0, + label: buildModelLabel(m), + size: m.size_gb || m.size_gb_approx || 0, category: m.category_expected || '', + cached: m.cached || false, + isMoe: m.is_moe || false, + fitsRam: m.fits_ram !== false, + notes: m.notes || '', })); } catch(e) { - // Fallback to a minimal list if registry fails KNOWN_MODELS = [ - {id:'mlx-community/Llama-3.2-3B-Instruct-4bit', label:'Llama 3.2 3B (2GB)', size:2, category:'small_dense'}, - {id:'mlx-community/Qwen3-30B-A3B-4bit', label:'Qwen3 30B MoE (18GB)', size:18, category:'large_moe'}, - {id:'mlx-community/Qwen3.6-35B-A3B-4bit', label:'Qwen3.6 35B MoE (20GB)', size:20, category:'large_moe'}, + {id:'mlx-community/Llama-3.2-3B-Instruct-4bit', label:'Llama 3.2 3B (2GB)', size:2, category:'small_dense', cached:false, isMoe:false, fitsRam:true}, ]; } - // Also add the currently active model if not in list try { const st = await fetch('/status'); const data = await st.json(); if (data.model) { activeModelId = data.model; if (!KNOWN_MODELS.some(m => m.id === activeModelId)) { - KNOWN_MODELS.unshift({id: activeModelId, label: activeModelId.split('/').pop() + ' (active)', size: 0, category: ''}); + KNOWN_MODELS.unshift({id: activeModelId, label: activeModelId.split('/').pop() + ' (active)', size: 0, category: '', cached: true, isMoe: false, fitsRam: true}); } } } catch(e) {} populateModelSelect(activeModelId || localStorage.getItem('mlx-flash-model') || ''); } +function buildModelLabel(m) { + const name = (m.id || m.name || '').split('/').pop(); + const size = m.size_gb || m.size_gb_approx || 0; + let label = name; + if (size > 0) label += ' (' + size + 'GB)'; + if (m.cached) label += ' \u2713'; + if (m.is_moe) label += ' [MoE]'; + return label; +} + const CATEGORY_LABELS = { - 'small_dense': 'Small Dense', - 'medium_dense': 'Medium Dense', - 'large_dense': 'Large Dense', + 'small_dense': 'Small Dense (\u2264 5GB)', + 'medium_dense': 'Medium Dense (5-20GB)', + 'large_dense': 'Large Dense (20GB+)', 'small_moe': 'Small MoE', 'medium_moe': 'Medium MoE', 'large_moe': 'Large MoE', + 'ssd_small': 'Hybrid SSM+Attn', '': 'Other', }; function populateModelSelect(currentModel) { const sel = document.getElementById('model-select'); - // Group models by category const groups = {}; const isKnown = KNOWN_MODELS.some(m => m.id === currentModel); - // If the active model isn't in the registry, add it at the top let topOpts = ''; if (currentModel && !isKnown) { const short = currentModel.split('/').pop() || currentModel; @@ -696,8 +707,7 @@ function populateModelSelect(currentModel) { groups[cat].push(m); }); let html = topOpts; - // Sort categories: small_dense, medium_dense, large_dense, small_moe, medium_moe, large_moe, then other - const catOrder = ['small_dense','medium_dense','large_dense','small_moe','medium_moe','large_moe','']; + const catOrder = ['small_dense','medium_dense','large_dense','small_moe','medium_moe','ssd_small','large_moe','']; const sortedCats = Object.keys(groups).sort((a,b) => { const ai = catOrder.indexOf(a), bi = catOrder.indexOf(b); return (ai===-1?99:ai) - (bi===-1?99:bi); @@ -708,31 +718,64 @@ function populateModelSelect(currentModel) { html += ''; groups[cat].forEach(m => { const isActive = m.id === currentModel; - const displayLabel = isActive ? m.label + ' *' : m.label; + let displayLabel = isActive ? '\u25B6 ' + m.label : m.label; html += ''; }); html += ''; }); } else { - // Flat list if all same category KNOWN_MODELS.forEach(m => { const isActive = m.id === currentModel; - const displayLabel = isActive ? m.label + ' *' : m.label; + let displayLabel = isActive ? '\u25B6 ' + m.label : m.label; html += ''; }); } - html += ''; + html += ''; sel.innerHTML = html; } +let switchPollInterval = null; + async function switchModel(modelId) { if (modelId === '_custom') { const custom = prompt('Enter HuggingFace model ID (e.g., mlx-community/gemma-4-31b-it-4bit):'); - if (!custom) return; + if (!custom) { populateModelSelect(activeModelId); return; } modelId = custom; } const shortName = modelId.split('/').pop(); - document.getElementById('status-text').textContent = 'Switching to ' + shortName + '...'; + const modelInfo = KNOWN_MODELS.find(m => m.id === modelId); + const isCached = modelInfo ? modelInfo.cached : false; + const sizeGb = modelInfo ? modelInfo.size : 0; + + // Show initial progress + const statusEl = document.getElementById('status-text'); + const dotEl = document.getElementById('status-dot'); + dotEl.style.background = 'var(--yellow)'; + if (!isCached && sizeGb > 0) { + statusEl.textContent = 'Downloading ' + shortName + ' (~' + sizeGb + 'GB)...'; + } else { + statusEl.textContent = 'Loading ' + shortName + '...'; + } + + // Poll progress while switch is happening + switchPollInterval = setInterval(async () => { + try { + const pr = await fetch('/switch/progress').then(r=>r.json()); + if (pr.switching) { + const phase = pr.phase || 'loading'; + if (phase === 'downloading') { + statusEl.textContent = 'Downloading ' + shortName + '...'; + } else if (phase === 'loading') { + statusEl.textContent = 'Loading ' + shortName + ' into GPU memory...'; + } else if (phase === 'warming') { + statusEl.textContent = 'Compiling Metal shaders for ' + shortName + '...'; + } else { + statusEl.textContent = phase + ' ' + shortName + '...'; + } + } + } catch(e) {} + }, 800); + try { const resp = await fetch('/v1/models/switch', { method: 'POST', @@ -740,18 +783,38 @@ async function switchModel(modelId) { body: JSON.stringify({model: modelId}), }); const result = await resp.json(); + clearInterval(switchPollInterval); + switchPollInterval = null; + if (result.switched) { - document.getElementById('status-text').textContent = 'Switched to ' + shortName; + dotEl.style.background = 'var(--green)'; + statusEl.textContent = 'Switched to ' + shortName; localStorage.setItem('mlx-flash-model', modelId); - } else if (result.error) { - document.getElementById('status-text').textContent = 'Switch failed: ' + result.error; + activeModelId = modelId; + // Refresh registry to update cached status + loadModelRegistry(); } else { - localStorage.setItem('mlx-flash-model', modelId); - document.getElementById('status-text').textContent = 'Model set: ' + shortName + ' (restart to apply)'; + dotEl.style.background = 'var(--red)'; + // Parse error from nested failures + let errMsg = result.error || ''; + if (!errMsg && result.failures && result.failures.length > 0) { + try { + const inner = JSON.parse(result.failures[0].error || '{}'); + errMsg = inner.error || result.failures[0].error || 'Unknown error'; + } catch(e) { errMsg = result.failures[0].error || 'Unknown error'; } + } + statusEl.textContent = 'Failed: ' + errMsg; + // Restore select to current model + populateModelSelect(activeModelId); + setTimeout(() => { dotEl.style.background = 'var(--green)'; statusEl.textContent = 'Ready'; }, 8000); } } catch(e) { - localStorage.setItem('mlx-flash-model', modelId); - document.getElementById('status-text').textContent = 'Model set: ' + shortName + ' (restart to apply)'; + clearInterval(switchPollInterval); + switchPollInterval = null; + dotEl.style.background = 'var(--red)'; + statusEl.textContent = 'Switch failed: ' + e.message; + populateModelSelect(activeModelId); + setTimeout(() => { dotEl.style.background = 'var(--green)'; statusEl.textContent = 'Ready'; }, 5000); } } diff --git a/mlx-flash-server/src/server.rs b/mlx-flash-server/src/server.rs index c084a69..857c707 100644 --- a/mlx-flash-server/src/server.rs +++ b/mlx-flash-server/src/server.rs @@ -70,6 +70,7 @@ pub fn create_router(state: AppState) -> Router { .route("/telemetry/current", get(handle_telemetry_current)) .route("/commands", get(handle_commands_list)) .route("/commands/run", axum::routing::post(handle_command_run)) + .route("/switch/progress", get(handle_switch_progress)) .route("/v1/models/registry", get(handle_model_registry)) .route("/v1/models/registry/cleanup", axum::routing::post(handle_model_cleanup)) .route("/v1/models/profile", axum::routing::post(handle_model_profile)) @@ -211,6 +212,24 @@ async fn handle_telemetry_current(State(state): State) -> axum::Json) -> axum::Json { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(3)) + .build() + .unwrap(); + + for port in state.pool.ports() { + let url = format!("http://127.0.0.1:{port}/switch/progress"); + if let Ok(resp) = client.get(&url).send().await { + if let Ok(body) = resp.json::().await { + return axum::Json(body); + } + } + } + + axum::Json(json!({"switching": false, "error": "No worker available"})) +} + async fn handle_release(State(state): State) -> axum::Json { state.request_count.fetch_add(1, Ordering::Relaxed); @@ -810,6 +829,26 @@ async fn handle_model_registry(State(state): State) -> axum::Json(); + let mib = [libc::CTL_HW, libc::HW_MEMSIZE]; + unsafe { + libc::sysctl( + mib.as_ptr() as *mut _, + 2, + &mut size as *mut _ as *mut _, + &mut len, + std::ptr::null_mut(), + 0, + ); + } + size as f64 / 1073741824.0 + }; + // Usable budget: ~75% of RAM (OS + other apps take ~25%) + let usable_gb = total_ram_gb * 0.75; + let mut models = Vec::new(); if let Some(model_list) = registry["models"].as_array() { for m in model_list { @@ -823,12 +862,23 @@ async fn handle_model_registry(State(state): State) -> axum::Json) -> axum::Json 80 tok/s, DFlash: skip) --- + # --- Small dense (< 5GB, fast AR) --- - id: mlx-community/Llama-3.2-3B-Instruct-4bit category_expected: small_dense size_gb_approx: 2 notes: "Llama 3.2 3B, dense, 4-bit" - - id: mlx-community/gemma-4-e2b-it-4bit + - id: mlx-community/Llama-3.2-1B-Instruct-4bit category_expected: small_dense size_gb_approx: 1 - notes: "Gemma 4 E2B (1.2B), dense, 4-bit — may fail (unsupported arch)" - skip: true + notes: "Llama 3.2 1B, dense, 4-bit — fastest" + + - id: mlx-community/Qwen3-4B-4bit + category_expected: small_dense + size_gb_approx: 3 + notes: "Qwen 3 4B, dense, 4-bit" + + - id: mlx-community/Phi-4-mini-instruct-4bit + category_expected: small_dense + size_gb_approx: 2 + notes: "Phi 4 Mini 3.8B, dense, 4-bit (Microsoft)" + + # --- Medium dense (5-20GB) --- + - id: mlx-community/Qwen3-8B-4bit + category_expected: medium_dense + size_gb_approx: 5 + notes: "Qwen 3 8B, dense, 4-bit" + + - id: mlx-community/Llama-3.3-70B-Instruct-4bit + category_expected: large_dense + size_gb_approx: 40 + notes: "Llama 3.3 70B, dense, 4-bit — needs 48GB+" - # --- Medium dense (AR 15-40 tok/s, DFlash: medium_target) --- - id: mlx-community/Qwen3.5-27B-4bit category_expected: medium_dense size_gb_approx: 15 @@ -31,25 +46,23 @@ models: size_gb_approx: 14 notes: "Mistral Devstral 24B, dense, 4-bit (code-focused)" - # --- Large dense (AR < 15 tok/s, DFlash: slow_target) --- + - id: mlx-community/gemma-4-12b-it-4bit + category_expected: medium_dense + size_gb_approx: 7 + notes: "Gemma 4 12B, dense, 4-bit (Google)" + + # --- Large dense (20GB+) --- - id: mlx-community/gemma-4-31b-it-4bit category_expected: large_dense size_gb_approx: 18 - notes: "Gemma 4 31B, dense, 4-bit" - - # --- Small MoE (AR > 50 tok/s, DFlash: skip) --- - - id: mlx-community/Qwen3.6-35B-A3B-4bit - category_expected: ssd_small - size_gb_approx: 19 - drafter: z-lab/Qwen3.6-35B-A3B-DFlash - notes: "Qwen 3.6 35B MoE (3B active), 30 SSM + 10 attn, 4-bit" + notes: "Gemma 4 31B, dense, 4-bit (Google)" + # --- Small MoE --- - id: mlx-community/Qwen3-30B-A3B-4bit category_expected: small_moe size_gb_approx: 16 notes: "Qwen 3 30B MoE (3B active), 4-bit" - # --- Medium MoE --- - id: mlx-community/Qwen3.5-35B-A3B-4bit category_expected: small_moe size_gb_approx: 20 @@ -58,11 +71,16 @@ models: - id: mlx-community/gemma-4-26b-a4b-it-4bit category_expected: small_moe size_gb_approx: 14 - notes: "Gemma 4 26B MoE (4B active), 4-bit" - - # --- Large MoE (AR < 15 tok/s, DFlash: slow_target) --- - - id: mlx-community/DeepSeek-V4-Flash-2bit-DQ - category_expected: large_moe - size_gb_approx: 90 - notes: "DeepSeek V4 Flash 284B (13B active), 2-bit DQ. Requires 128GB+ or expert streaming." - skip_on_small_ram: true + notes: "Gemma 4 26B MoE (4B active), 4-bit (Google)" + + # --- Medium/Large MoE --- + - id: mlx-community/Qwen3.6-35B-A3B-4bit + category_expected: medium_moe + size_gb_approx: 19 + drafter: z-lab/Qwen3.6-35B-A3B-DFlash + notes: "Qwen 3.6 35B MoE (3B active), SSM+attn hybrid, 4-bit" + + - id: mlx-community/Mixtral-8x7B-Instruct-v0.1-4bit + category_expected: medium_moe + size_gb_approx: 24 + notes: "Mixtral 8x7B MoE (12B active), 4-bit (Mistral)" From 967c21eae63d3b82d908b14296c9b01e6a680e2b Mon Sep 17 00:00:00 2001 From: szibis <329831+szibis@users.noreply.github.com> Date: Mon, 11 May 2026 15:07:39 +0200 Subject: [PATCH 5/6] =?UTF-8?q?fix:=20SSE=20streaming=20double-wrap=20?= =?UTF-8?q?=E2=80=94=20pass=20through=20raw=20bytes=20from=20Python=20work?= =?UTF-8?q?er?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rust proxy was wrapping Python's "data: {...}" lines in another Event::data(), producing "data: data: {...}" which broke the chat UI's SSE parser. Now passes raw byte stream through with correct headers. --- mlx-flash-server/src/proxy.rs | 29 +++++++++++------------------ 1 file changed, 11 insertions(+), 18 deletions(-) diff --git a/mlx-flash-server/src/proxy.rs b/mlx-flash-server/src/proxy.rs index 813ccd0..44c40c2 100644 --- a/mlx-flash-server/src/proxy.rs +++ b/mlx-flash-server/src/proxy.rs @@ -1,9 +1,7 @@ use std::sync::atomic::Ordering; use axum::extract::State; -use axum::response::{IntoResponse, Response}; -use axum::response::sse::{Event, KeepAlive, Sse}; -use tokio_stream::StreamExt; +use axum::response::Response; use crate::server::AppState; @@ -107,22 +105,17 @@ pub async fn handle_chat( } Ok(upstream) => { if streaming { + // Pass through raw SSE bytes from Python — it already sends + // properly formatted "data: {...}\n\n" lines let byte_stream = upstream.bytes_stream(); - let event_stream = byte_stream.map(|chunk| { - chunk - .map(|bytes| { - Event::default().data( - String::from_utf8_lossy(&bytes) - .trim_end() - .to_string(), - ) - }) - .map_err(|e| format!("stream error: {e}")) - }); - - Sse::new(event_stream) - .keep_alive(KeepAlive::default()) - .into_response() + let body = axum::body::Body::from_stream(byte_stream); + axum::response::Response::builder() + .status(200) + .header("Content-Type", "text/event-stream") + .header("Cache-Control", "no-cache") + .header("Connection", "keep-alive") + .body(body) + .unwrap() } else { let status = upstream.status(); let bytes = match upstream.bytes().await { From 92d25d47bb02b9a14aa6552667fd1f9132017683 Mon Sep 17 00:00:00 2001 From: szibis <329831+szibis@users.noreply.github.com> Date: Mon, 11 May 2026 15:51:22 +0200 Subject: [PATCH 6/6] fix: live token counter, per-message stats, session persistence across nav - Token counter updates live during SSE streaming (counts content chunks) - Status bar shows tokens, tok/s, elapsed time after each message - Session saved on beforeunload (covers refresh, nav to dashboard, tab close) - loadSession validates data before restoring --- mlx-flash-server/src/chat_ui.rs | 68 ++++++++++++++++++--------------- 1 file changed, 38 insertions(+), 30 deletions(-) diff --git a/mlx-flash-server/src/chat_ui.rs b/mlx-flash-server/src/chat_ui.rs index c937e3a..84d1ebb 100644 --- a/mlx-flash-server/src/chat_ui.rs +++ b/mlx-flash-server/src/chat_ui.rs @@ -306,31 +306,31 @@ function formatContent(text) { } let genPollInterval = null; -let tokBefore = 0; +let liveTokenCount = 0; function startGenProgress() { const el = document.getElementById('gen-progress'); el.classList.add('active'); document.getElementById('status-dot').style.background = 'var(--accent)'; - tokBefore = 0; - // Snapshot current tokens - fetch('/status').then(r=>r.json()).then(st => { tokBefore = (st.stats||{}).tokens_generated||0; }).catch(()=>{}); - // Poll every 500ms during generation + liveTokenCount = 0; + document.getElementById('gen-tokens').textContent = '0 tok'; + // Poll memory/pressure during generation genPollInterval = setInterval(async () => { try { const st = await fetch('/status').then(r=>r.json()); - const tokNow = (st.stats||{}).tokens_generated||0; - const genTok = tokNow - tokBefore; - document.getElementById('gen-tokens').textContent = genTok + ' tok'; const mem = st.memory||{}; const avail = ((mem.free_gb||0)+(mem.inactive_gb||0)*0.5).toFixed(1); const pressure = (mem.pressure||'Normal').toString(); document.getElementById('gen-mem').textContent = avail+'GB free'; document.getElementById('gen-mem').style.color = pressure==='Critical'?'var(--red)':pressure==='Warning'?'var(--yellow)':'var(--dim)'; - // Pressure banner updatePressureBanner(pressure, avail, mem.swap_used_gb||0); } catch(e) {} - }, 500); + }, 1000); +} + +function updateLiveTokens(count) { + liveTokenCount = count; + document.getElementById('gen-tokens').textContent = count + ' tok'; } function stopGenProgress() { @@ -357,22 +357,30 @@ function updatePressureBanner(pressure, availGb, swapGb) { } } -// Session persistence — survive page reload / tab switch +// Session persistence — survives page reload, tab switch, nav to dashboard and back function saveSession() { - try { localStorage.setItem('mlx-flash-messages', JSON.stringify(messages)); } catch(e) {} + try { + localStorage.setItem('mlx-flash-messages', JSON.stringify(messages)); + localStorage.setItem('mlx-flash-session-ts', Date.now().toString()); + } catch(e) {} } function loadSession() { try { const saved = localStorage.getItem('mlx-flash-messages'); - if (saved) { - const restored = JSON.parse(saved); - if (restored.length > 0) { - document.getElementById('empty')?.remove(); - restored.forEach(m => addMessage(m.role, m.content)); + if (!saved) return; + const restored = JSON.parse(saved); + if (!Array.isArray(restored) || restored.length === 0) return; + document.getElementById('empty')?.remove(); + restored.forEach(m => { + if (m && m.role && m.content) { + addMessage(m.role, m.content); } - } - } catch(e) {} + }); + } catch(e) { console.error('Session restore failed:', e); } } +// Save on every page unload (covers nav away, refresh, tab close) +window.addEventListener('beforeunload', saveSession); +// Restore on load loadSession(); async function handleSlashCommand(text) { @@ -524,12 +532,12 @@ async function sendMessage() { // SSE streaming — read tokens as they arrive let fullText = ''; let respTokens = 0; + let streamChunkCount = 0; let tokPerSec = ''; let lastChunkModel = currentModel; const contentType = resp.headers.get('content-type') || ''; if (contentType.includes('text/event-stream') || contentType.includes('text/plain')) { - // SSE stream mode const reader = resp.body.getReader(); const decoder = new TextDecoder(); let buffer = ''; @@ -550,22 +558,25 @@ async function sendMessage() { const delta = chunk.choices?.[0]?.delta || {}; if (delta.content) { fullText += delta.content; + streamChunkCount++; + updateLiveTokens(streamChunkCount); aiEl.innerHTML = formatContent(fullText); messagesEl.scrollTop = messagesEl.scrollHeight; } if (chunk.model) lastChunkModel = chunk.model; - // Some implementations send usage in the final chunk if (chunk.usage) { respTokens = chunk.usage.completion_tokens || 0; } - // Extract MLX-Flash performance data if present const mlxData = chunk.mlx_flash_compress || {}; if (mlxData.tok_per_s) tokPerSec = mlxData.tok_per_s.toFixed(1) + ' tok/s'; - } catch(e) { /* skip unparseable lines */ } + } catch(e) {} } } + // Use chunk count as token estimate if usage not provided + if (!respTokens && streamChunkCount > 0) { + respTokens = streamChunkCount; + } } else { - // Fallback: non-streaming JSON response try { const json = await resp.json(); fullText = json.choices?.[0]?.message?.content @@ -583,7 +594,6 @@ async function sendMessage() { } const elapsed = ((Date.now() - t0) / 1000).toFixed(1); - // Estimate tokens from text if not provided if (!respTokens && fullText) { respTokens = Math.ceil(fullText.length / 4); } @@ -592,14 +602,13 @@ async function sendMessage() { } messages[messages.length-1].content = fullText; - // Track cumulative tokens for savings calculator if (respTokens > 0) { totalTokensSession += respTokens; try { localStorage.setItem('mlx-flash-total-tokens', totalTokensSession.toString()); } catch(e) {} } saveSession(); - // Get current model + pressure for metadata bar + // Per-message stats let currentPressure = 'Normal'; let currentModelName = lastChunkModel || currentModel; try { @@ -608,7 +617,6 @@ async function sendMessage() { if (st.model) currentModelName = st.model; } catch(e) {} - // Set metadata bar under the AI message setMessageMeta(messages.length - 1, { model: currentModelName, tokens: respTokens || null, @@ -617,11 +625,11 @@ async function sendMessage() { pressure: currentPressure.toString(), }); - // Show generation stats in status bar const stats = [elapsed + 's']; + if (respTokens) stats.push(respTokens + ' tok'); if (tokPerSec) stats.push(tokPerSec); document.getElementById('status-text').textContent = 'Done — ' + stats.join(', '); - setTimeout(() => { if (!generating) document.getElementById('status-text').textContent = 'Ready'; }, 5000); + setTimeout(() => { if (!generating) document.getElementById('status-text').textContent = 'Ready'; }, 8000); } catch(e) { aiEl.innerHTML = '

Connection error — is the Python worker running?

';