mtplx 2.0.0__tar.gz → 2.0.2__tar.gz

{mtplx-2.0.0 → mtplx-2.0.2}/CHANGELOG.md

@@ -4,6 +4,142 @@ All notable user-facing changes to MTPLX. The format is based on
 [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versions follow
 [Semantic Versioning](https://semver.org/).
 
+## [2.0.2] - 2026-07-09
+
+The agent-reliability release: multi-turn agent sessions now render
+reasoning history on the model's trained contract (the fix for the
+plan-execution repetition marathons), LAN serving works from the app,
+and every agent client gets warm prefix reuse.
+
+### Changed
+
+- Reasoning history is now scoped to the active agent round by default
+  on Qwen3.6/3.5 models, matching Qwen's trained multi-turn contract.
+  The Qwen chat template keeps `<think>` blocks only for assistant
+  messages after the last real user query (its built-in "rolling
+  checkpoint"); MTPLX used to override that with `preserve_thinking`,
+  which rendered an off-contract empty `<think>` scaffold on every
+  completed assistant turn and replayed stale inline reasoning across
+  turns, while silently dropping the structured `reasoning_content`
+  fields agent clients such as OpenCode send. Scoped mode lets the
+  template's own checkpoint govern: completed turns render with no
+  think scaffold, and the active round (assistant -> tool -> assistant
+  chains) keeps its reasoning, now including structured
+  `reasoning_content` - strictly better in-round continuity than
+  before. `--preserve-thinking on` restores the previous behavior
+  byte-for-byte (including cache identity), `off` still strips
+  everything, and the new explicit `scoped` value pins the scoped mode.
+  Templates without the rolling checkpoint (Gemma 4, custom templates)
+  keep the previous preserve-all behavior. The resolved policy is shown
+  at startup ("Reasoning history: scoped (active round only)") and as
+  `reasoning_history_mode` in `/v1/mtplx/settings` and the snapshot.
+
+### Added
+
+- Loop Guard (opt-in): `MTPLX_LOOP_GUARD=1` enables a loop-armed
+  anti-repetition steering mode for models prone to verbatim repetition
+  marathons (for example repetition-damaged third-party quants). Unlike
+  a static presence penalty, the guard is completely inert until a real
+  loop is detected mid-response (bit-exact sampling otherwise, MTP
+  acceptance math unchanged), then penalizes only the tokens that would
+  extend a verbatim repeat and disarms once the loop is broken. Content
+  inside tool calls is never steered (code legitimately repeats short
+  token runs, so tool-call spans are masked token-exactly). The default
+  is OFF: no synthetic steering touches sampling unless you ask for it,
+  and the repetition fix that matters for MTPLX's own models is the
+  scoped reasoning history change above. Detector/steering knobs are
+  documented in `mtplx/loop_guard.py`; per-request guard activity is
+  visible under `loop_guard` in `/v1/mtplx/snapshot`.
+
+### Fixed
+
+- Hidden reasoning no longer leaks into visible chat content when a
+  long reasoning tag (e.g. `</reasoning>`) splits across a streaming
+  chunk boundary; the splitter held back too few bytes for tags longer
+  than `</think>` (PR #149 by @Osamaali313).
+- `mtplx` quickstart onboarding screens now display the requested
+  `--host`/`--port` in the Web UI and dashboard URLs instead of a
+  hardcoded `http://127.0.0.1:8000` (PR #148 by @hasegaw).
+- Ctrl-C now returns control to the terminal within a bounded delay
+  even when a browser tab holds an open chat/dashboard stream. The
+  server previously waited forever for infinite SSE generators to
+  finish; shutdown now drains in-flight requests with a 5-second
+  deadline, and thermal/fan cleanup still runs (#124).
+- Serving on any host and port now works from the app. Setting host
+  0.0.0.0 (LAN serving) used to misreport free ports as occupied, bump
+  the port, and kill the healthy daemon after a health-wait timeout,
+  because the app probed the bind address verbatim; all app-side
+  connections now resolve wildcard binds to a connectable loopback
+  address (#109). The app also surfaces the "LAN serving requires an
+  API key" rule before launch instead of a generic Degraded state, an
+  API-key mismatch reads as a live-but-unauthorized daemon instead of
+  a lost one, and the port-in-use preflight tests the address family
+  the daemon will actually bind.
+- The app now passes the SSD session cache setting to the daemon
+  explicitly, including "Off". Since 2.0.0 flipped the serve default to
+  on, an explicit Settings "Off" was silently re-enabled with default
+  limits on app-launched daemons, and `~/.mtplx/session-bank` kept
+  growing (#140; also the "session-bank came back after I deleted it"
+  half of #138). Generated `mtplx start` server commands carry the
+  explicit `--ssd-session-cache off` for the same reason.
+- The app's runtime venv now self-heals after app updates. A venv whose
+  base-interpreter symlink pointed into a replaced app bundle made every
+  reinstall fail with "[Errno 2] No such file or directory:
+  .../runtime-venv/bin/python3" and no DMG reinstall could fix it; venv
+  creation now rebuilds with `--clear` when the existing venv python is
+  broken or creation fails (#139).
+- Warm prefix reuse no longer freezes on the oldest short prefix for
+  agent harnesses outside OpenCode (Pi, little-coder, Hermes, custom
+  clients). The block-prefix restore lane was still gated to OpenCode's
+  compact tool contract even though kvcache-v2's boundary-true restores
+  made it exact for every client, so transcripts whose turns diverge
+  more than a few tokens before the stored end re-prefilled a growing
+  suffix every turn (#138). The lane now engages for all clients while
+  boundary-true restore is on; `MTPLX_SESSION_BLOCK_PREFIX_RESTORE=0`
+  still disables it.
+
+## [2.0.1] - 2026-07-07
+
+Turbo for every Mac. The v2 turbo default now covers every dense catalog
+model on every Apple Silicon generation, with a load-time kernel
+self-validation safety net.
+
+### Added
+
+- 6-bit affine verify kernels (split-K hexpack family): the 6-bit 9B tier
+  gains 33-62% decode and 43% 2k-prefill under turbo (M5 Max, verified
+  arms). Qwen 3.5 9B and 9B FP16 now default to turbo.
+- New model: `Youssofal/Qwen3.6-27B-MTPLX-Optimized-Quality-FP16`, the
+  missing M1/M2 quality artifact. Wired into the app picker, CLI catalog,
+  and chip-aware routing (a Quality pick on M1/M2 resolves the FP16
+  sibling, mirroring the speed lane). Validated against the bf16 parent
+  and measured at 2.5x over true AR under turbo.
+- Load-time kernel self-validation: every turbo lane checks itself against
+  stock MLX in the model's exact dtype/quantization at boot; a mismatching
+  lane falls back to the stock path for the session and the verdicts are
+  surfaced as `kernel_selfcheck` in `/health`. Worst case is v2.0.0 speed,
+  never wrong output.
+- `MTPLX_FORCE_GPU_FAMILY_FALLBACK=1` rehearses the exact M1-M4 code path
+  on newer machines; `MTPLX_KERNEL_SELFCHECK=0` disables the probe.
+- CI kernel matrix on real M1 runners: kernel exactness across
+  {bf16, fp16} x {4, 6, 8}-bit plus a live 4B turbo boot smoke.
+
+### Changed
+
+- 27B Optimized-Speed-FP16 (the M1/M2 routing target) defaults to turbo:
+  19-31% faster decode than the v2.0.0 sustained default across 0.5k-32k
+  context; true-AR multiplier 1.34x to ~2x. First-ever e2e measurement of
+  this artifact; exactness gated (30/30 hot-shape logit-diff cases on real
+  weights, greedy match turbo vs stock).
+- The published Speed-FP16 model card now recommends the turbo profile.
+
+### Unchanged on purpose
+
+- 35B A3B MoE and Gemma 4 keep sustained (expert layers / assistant-pair
+  architecture bypass these kernels; named 2.0.2 lanes). The 4B keeps
+  sustained (turbo measured slightly slower at matched depth). Compiled
+  verify stays off for 6-bit models.
+
 ## [2.0.0] - 2026-07-06
 
 MTPLX v2: the coding-agent release. Session-cache v2 (RAM + SSD), the

{mtplx-2.0.0 → mtplx-2.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mtplx
-Version: 2.0.0
+Version: 2.0.2
 Summary: Native MTP speculative decoding for Qwen3-Next on Apple Silicon.
 Author-email: Youssof Altoukhi <business@youssofal.com>
 License-Expression: Apache-2.0

{mtplx-2.0.0 → mtplx-2.0.2}/mtplx/artifacts.py

@@ -32,6 +32,8 @@ from .profiles import (
     DEFAULT_PUBLIC_MODEL_ID,
     LEGACY_OPTIMIZED_HF_MODEL_ID,
     LEGACY_OPTIMIZED_PUBLIC_MODEL_ID,
+    QUALITY_FP16_HF_MODEL_ID,
+    QUALITY_FP16_PUBLIC_MODEL_ID,
     QUALITY_HF_MODEL_ID,
     QUALITY_PUBLIC_MODEL_ID,
     QWEN35_9B_OPTIMIZED_SPEED_FP16_HF_MODEL_ID,
@@ -57,6 +59,7 @@ _KNOWN_PUBLIC_MODEL_ALIASES = {
     DEFAULT_PUBLIC_MODEL_ID: DEFAULT_HF_MODEL_ID,
     DEFAULT_FP16_PUBLIC_MODEL_ID: DEFAULT_FP16_HF_MODEL_ID,
     QUALITY_PUBLIC_MODEL_ID: QUALITY_HF_MODEL_ID,
+    QUALITY_FP16_PUBLIC_MODEL_ID: QUALITY_FP16_HF_MODEL_ID,
     LEGACY_OPTIMIZED_PUBLIC_MODEL_ID: LEGACY_OPTIMIZED_HF_MODEL_ID,
     QWEN35_9B_OPTIMIZED_SPEED_PUBLIC_MODEL_ID: QWEN35_9B_OPTIMIZED_SPEED_HF_MODEL_ID,
     QWEN35_9B_OPTIMIZED_SPEED_FP16_PUBLIC_MODEL_ID: QWEN35_9B_OPTIMIZED_SPEED_FP16_HF_MODEL_ID,
@@ -71,6 +74,7 @@ _KNOWN_PUBLIC_MODEL_ALIASES = {
     "qwen3.6-27b-mtplx-optimized": LEGACY_OPTIMIZED_HF_MODEL_ID,
     "qwen3.6-27b-mtplx-optimized-speed-fp16": DEFAULT_FP16_HF_MODEL_ID,
     "qwen3.6-27b-mtplx-optimized-quality": QUALITY_HF_MODEL_ID,
+    "qwen3.6-27b-mtplx-optimized-quality-fp16": QUALITY_FP16_HF_MODEL_ID,
     "qwen3.6-35b-a3b-mtplx-optimized-speed": QWEN36_35B_OPTIMIZED_SPEED_HF_MODEL_ID,
     "qwen3.6-35b-a3b-mtplx-optimized-speed-fp16": QWEN36_35B_OPTIMIZED_SPEED_FP16_HF_MODEL_ID,
     "qwen3.6-35b-a3b-mtplx-optimized-balance": QWEN36_35B_OPTIMIZED_BALANCE_HF_MODEL_ID,

{mtplx-2.0.0 → mtplx-2.0.2}/mtplx/attention_split.py

@@ -234,6 +234,10 @@ def _install_split_attention_hook(attn: Any) -> bool:
         gqa_packed_threshold = int(
             getattr(self, "_mtplx_gqa_packed_sdpa_threshold", 8192)
         )
+        if gqa_packed_enabled:
+            from .kernel_selfcheck import lane_disabled
+
+            gqa_packed_enabled = not lane_disabled("gqa_packed_sdpa")
         should_use_gqa_packed = (
             gqa_packed_enabled
             and cache is not None

{mtplx-2.0.0 → mtplx-2.0.2}/mtplx/cli.py

@@ -542,11 +542,13 @@ def _add_reasoning_effort_arg(parser: argparse.ArgumentParser) -> None:
 def _add_preserve_thinking_arg(parser: argparse.ArgumentParser) -> None:
     parser.add_argument(
         "--preserve-thinking",
-        choices=["auto", "on", "off"],
+        choices=["auto", "on", "off", "scoped"],
         default="auto",
         help=(
-            "Preserve prior assistant reasoning in Qwen chat-template history. "
-            "Default auto preserves it for reasoning-capable templates; off is a speed/debug mode."
+            "Reasoning-history policy for Qwen chat-template history. scoped keeps "
+            "reasoning only inside the active agent round (Qwen's trained contract); "
+            "on preserves all; off strips all. Default auto resolves to scoped for "
+            "checkpoint-capable templates."
         ),
     )
     parser.add_argument(

{mtplx-2.0.0 → mtplx-2.0.2}/mtplx/commands/public.py

@@ -83,6 +83,8 @@ from mtplx.profiles import (
     DEFAULT_PUBLIC_MODEL_ID,
     LEGACY_OPTIMIZED_HF_MODEL_ID,
     LEGACY_OPTIMIZED_PUBLIC_MODEL_ID,
+    QUALITY_FP16_HF_MODEL_ID,
+    QUALITY_FP16_PUBLIC_MODEL_ID,
     QUALITY_HF_MODEL_ID,
     QUALITY_PUBLIC_MODEL_ID,
     QWEN35_9B_OPTIMIZED_SPEED_FP16_HF_MODEL_ID,
@@ -823,19 +825,40 @@ def _apply_model_contract_depth_default(
 
 # Quantized 27B flagships whose measured default profile is turbo. This is
 # the same launch rule the macOS app applies in MTPLXCommandBuilder
-# (Speed/Quality -> turbo; FP16 siblings stay sustained): NAX verify kernels
-# +22-40% chat decode on q8 (ULP-exact) and vk_k on q4, plus compiled verify
-# behind its per-model quant-bits gate. Before 2026-07-05 the bare CLI
-# (`mtplx serve` / quickstart / start) silently stayed on sustained, so every
-# OpenAI-API consumer — including third-party benchmark harnesses — measured
-# the slow path while the app ran turbo. Keep this list measured-win only:
-# 6-bit small models, 35B, Gemma, FP16 and third-party artifacts keep the
-# sustained default.
+# (Speed/Quality/Speed-FP16 -> turbo): NAX verify kernels +22-40% chat decode
+# on q8 (ULP-exact) and vk_k on q4, plus compiled verify behind its per-model
+# quant-bits gate. Before 2026-07-05 the bare CLI (`mtplx serve` / quickstart
+# / start) silently stayed on sustained, so every OpenAI-API consumer —
+# including third-party benchmark harnesses — measured the slow path while
+# the app ran turbo. Keep this list measured-win only: 6-bit small models
+# (vk covers 4/8-bit affine only), 35B MoE (experts bypass the NAX patch),
+# Gemma, and third-party artifacts keep the sustained default.
 _TURBO_DEFAULT_PUBLIC_MODEL_IDS = frozenset(
     {
         DEFAULT_PUBLIC_MODEL_ID,  # 27B Optimized-Speed (flat 4-bit)
         QUALITY_PUBLIC_MODEL_ID,  # 27B Optimized-Quality (8-bit)
         LEGACY_OPTIMIZED_PUBLIC_MODEL_ID,  # 27B Optimized (gdn8 hybrid, 8/4-bit)
+        # 27B Speed-FP16 (INT4/g64 weights, fp16 activations — the M1/M2
+        # routing target). Promoted 2026-07-07 after the first-ever e2e
+        # measurement: turbo 1.98-2.08x over true AR on M5 Max (55-60 tok/s
+        # D3 vs 27-29 AR), acceptance [86/66/51]%, peak 18.6 GB, vs only
+        # 1.34x on sustained (fp16 verify_hidden_eval tax ~8s/384-tok run).
+        # The 9B FP16 (6-bit) and 35B FP16 (MoE) siblings are NOT promoted.
+        DEFAULT_FP16_PUBLIC_MODEL_ID,
+        # 27B Quality-FP16 (q8/g64 weights, fp16 activations — the M1/M2
+        # quality pick, built 2026-07-07). Measured on its own artifact:
+        # turbo MTP D3 43.8/42.1 tok/s vs true AR 17.4 (2.5x), selfcheck
+        # all vk q8 lanes ok, parent logit-diff argmax 16/16.
+        QUALITY_FP16_PUBLIC_MODEL_ID,
+        # 9B (6-bit/g64) promoted 2026-07-07 after the 6-bit hexpack
+        # split-K kernels landed: live ABBA on the 9B measured turbo MTP
+        # D3 110/102 tok/s vs sustained 90/69 (+34%), true AR 62-65 flat
+        # both profiles -> multiplier 1.25x -> 1.68x. Exactness: 54
+        # synthetic {bf16,fp16}x{gs32,64,128} cases dmax <= 0.027 vs
+        # stock. The FP16 sibling shares the 6-bit packs and fp16 lanes
+        # are exactness-proven, so both promote together.
+        QWEN35_9B_OPTIMIZED_SPEED_PUBLIC_MODEL_ID,
+        QWEN35_9B_OPTIMIZED_SPEED_FP16_PUBLIC_MODEL_ID,
     }
 )
 
@@ -1162,7 +1185,7 @@ def _preserve_thinking_policy(args: Any) -> str:
         return "off"
     raw = getattr(args, "preserve_thinking", None)
     mode = str(raw or "auto").strip().lower()
-    return mode if mode in {"auto", "on", "off"} else "auto"
+    return mode if mode in {"auto", "on", "off", "scoped"} else "auto"
 
 
 def _pi_preserve_thinking_policy(args: Any) -> str:
@@ -7628,6 +7651,9 @@ def _model_ref_from_public_model_id(model_id: str | None) -> str | None:
         QUALITY_PUBLIC_MODEL_ID.lower(): QUALITY_HF_MODEL_ID,
         QUALITY_HF_MODEL_ID.lower(): QUALITY_HF_MODEL_ID,
         Path(QUALITY_HF_MODEL_ID).name.lower(): QUALITY_HF_MODEL_ID,
+        QUALITY_FP16_PUBLIC_MODEL_ID.lower(): QUALITY_FP16_HF_MODEL_ID,
+        QUALITY_FP16_HF_MODEL_ID.lower(): QUALITY_FP16_HF_MODEL_ID,
+        Path(QUALITY_FP16_HF_MODEL_ID).name.lower(): QUALITY_FP16_HF_MODEL_ID,
         QWEN35_9B_OPTIMIZED_SPEED_PUBLIC_MODEL_ID.lower(): QWEN35_9B_OPTIMIZED_SPEED_HF_MODEL_ID,
         QWEN35_9B_OPTIMIZED_SPEED_HF_MODEL_ID.lower(): QWEN35_9B_OPTIMIZED_SPEED_HF_MODEL_ID,
         Path(QWEN35_9B_OPTIMIZED_SPEED_HF_MODEL_ID).name.lower(): QWEN35_9B_OPTIMIZED_SPEED_HF_MODEL_ID,
@@ -9670,8 +9696,12 @@ def _batching_command_suffix(args: Any) -> str:
     if bool(getattr(args, "experimental_mtp_cohorts", False)):
         parts.append("--experimental-mtp-cohorts")
     ssd_session_cache = str(getattr(args, "ssd_session_cache", "on") or "on")
+    # Emit the mode unconditionally, "off" included: these generated
+    # commands are re-parsed by CLIs whose own default is "on"
+    # (kvcache-v2), so omitting an explicit "off" silently re-enables
+    # the cache (issue #140 class).
+    parts.extend(["--ssd-session-cache", shlex.quote(ssd_session_cache)])
     if ssd_session_cache != "off":
-        parts.extend(["--ssd-session-cache", shlex.quote(ssd_session_cache)])
         ssd_dir = getattr(args, "ssd_session_cache_dir", None)
         if ssd_dir:
             parts.extend(["--ssd-session-cache-dir", shlex.quote(str(ssd_dir))])
@@ -11956,6 +11986,8 @@ def cmd_quickstart_public(args: Any) -> int:
             fresh=fresh,
             configured_model=configured_model,
             open_dashboard_override=explicit_open_dashboard,
+            host=str(getattr(args, "host", "127.0.0.1")),
+            port=int(getattr(args, "port", 8000)),
         )
         if choice is None:
             _quickstart_line("aborted")

{mtplx-2.0.0 → mtplx-2.0.2}/mtplx/default_models.py

@@ -18,6 +18,8 @@ from mtplx.profiles import (
     DEFAULT_MODEL_ID,
     DEFAULT_PUBLIC_MODEL_ID,
     LEGACY_OPTIMIZED_PUBLIC_MODEL_ID,
+    QUALITY_FP16_HF_MODEL_ID,
+    QUALITY_FP16_PUBLIC_MODEL_ID,
     QUALITY_HF_MODEL_ID,
     QUALITY_PUBLIC_MODEL_ID,
     QWEN35_9B_OPTIMIZED_SPEED_FP16_HF_MODEL_ID,
@@ -60,6 +62,10 @@ _OPTIMIZED_QUALITY_LOCAL_CANDIDATES = (
     "~/Documents/MTPLX/hf-staging/Qwen3.6-27B-MTPLX-Optimized-Quality",
     "~/.mtplx/models/Youssofal--Qwen3.6-27B-MTPLX-Optimized-Quality",
 )
+_OPTIMIZED_QUALITY_FP16_LOCAL_CANDIDATES = (
+    "~/Documents/MTPLX/hf-staging/Qwen3.6-27B-MTPLX-Optimized-Quality-FP16",
+    "~/.mtplx/models/Youssofal--Qwen3.6-27B-MTPLX-Optimized-Quality-FP16",
+)
 _OPTIMIZED_35B_SPEED_LOCAL_CANDIDATES = (
     "~/Documents/MTPLX/models/Qwen3.6-35B-A3B-MTPLX-Optimized-Speed",
     "~/Documents/MTPLX/models/Qwen3.6-35B-A3B-MTPLX-Official4-CyanKiwiMTP-CleanRecipe",
@@ -171,16 +177,29 @@ def optimized_speed_model_ref() -> str:
     return local or DEFAULT_HF_MODEL_ID
 
 
-def optimized_quality_model_ref() -> str:
+def optimized_quality_model_ref(
+    *,
+    hardware: Mapping[str, Any] | None = None,
+) -> str:
     env_ref = str(os.environ.get(QUALITY_MODEL_ENV) or "").strip()
-    if env_ref:
-        if _env_ref_disabled(env_ref):
-            candidates = ()
-        else:
-            candidates = (env_ref, *_OPTIMIZED_QUALITY_LOCAL_CANDIDATES)
-    else:
-        candidates = _OPTIMIZED_QUALITY_LOCAL_CANDIDATES
-    local = _complete_local_model_ref(candidates)
+    env_disabled = _env_ref_disabled(env_ref) if env_ref else False
+    if env_ref and not env_disabled:
+        local = _complete_local_model_ref((env_ref,))
+        if local:
+            return local
+    # A quality pick on legacy (M1/M2) silicon resolves the FP16 sibling,
+    # mirroring the speed lane's precision routing (2.0.1, 2026-07-07).
+    hardware_info = dict(detect_apple_silicon() if hardware is None else hardware)
+    generation = _hardware_generation(hardware_info)
+    if generation in _LEGACY_APPLE_FP16_GENERATIONS:
+        if not env_disabled:
+            local = _complete_local_model_ref(_OPTIMIZED_QUALITY_FP16_LOCAL_CANDIDATES)
+            if local:
+                return local
+        return QUALITY_FP16_HF_MODEL_ID
+    if env_disabled:
+        return QUALITY_HF_MODEL_ID
+    local = _complete_local_model_ref(_OPTIMIZED_QUALITY_LOCAL_CANDIDATES)
     return local or QUALITY_HF_MODEL_ID
 
 
@@ -293,6 +312,8 @@ def _public_model_id_from_name(value: str) -> str | None:
         # First-party local research build of the released 35B speed
         # artifact (listed in _OPTIMIZED_35B_SPEED_LOCAL_CANDIDATES).
         return QWEN36_35B_OPTIMIZED_SPEED_PUBLIC_MODEL_ID
+    if "qwen3.6-27b-mtplx-optimized-quality-fp16" in lowered:
+        return QUALITY_FP16_PUBLIC_MODEL_ID
     if "qwen3.6-27b-mtplx-optimized-quality" in lowered:
         return QUALITY_PUBLIC_MODEL_ID
     if "qwen3.6-27b-mtplx-optimized-speed-fp16" in lowered:

{mtplx-2.0.0 → mtplx-2.0.2}/mtplx/fast_sampling.py

@@ -15,6 +15,7 @@ def apply_penalties_mlx(
     token_counts: Mapping[int, int] | None,
     presence_penalty: float = 0.0,
     frequency_penalty: float = 0.0,
+    penalty_overlay: Mapping[int, float] | None = None,
 ) -> mx.array:
     """On-device MLX twin of ``sampling.apply_penalties`` for one logit row.
 
@@ -22,20 +23,33 @@ def apply_penalties_mlx(
     temperature/top-k/top-p — via a sparse scatter over only the seen tokens
     (``logits.at[ids].add(-deltas)``): O(unique_seen), no dense vocab-sized
     allocation and no host round-trip beyond the small (ids, deltas) transfer.
-    Penalties clamp to [-2, 2]. Returns the same array unchanged (no-op) when
-    both penalties are 0 or there are no counts, preserving exactness.
+    Penalties clamp to [-2, 2]. ``penalty_overlay`` is the Loop Guard's sparse
+    token->subtraction map (not clamped; the guard caps it). Returns the same
+    array unchanged (no-op) when nothing is active, preserving exactness.
 
     Expects a 1-D logit row (``logits.shape == (vocab,)``); batched callers apply
     it per row (the MTP draft block is a handful of positions, not a hot loop).
     """
     presence = min(max(float(presence_penalty), PENALTY_MIN), PENALTY_MAX)
     frequency = min(max(float(frequency_penalty), PENALTY_MIN), PENALTY_MAX)
-    if (presence == 0.0 and frequency == 0.0) or not token_counts:
+    counts_active = bool(token_counts) and (presence != 0.0 or frequency != 0.0)
+    overlay_active = bool(penalty_overlay)
+    if not counts_active and not overlay_active:
         return logits
-    ids = np.fromiter(token_counts.keys(), dtype=np.int64, count=len(token_counts))
-    counts = np.fromiter(token_counts.values(), dtype=np.float64, count=len(token_counts))
-    deltas = (frequency * counts + presence * (counts > 0.0)).astype(np.float32)
-    return logits.at[mx.array(ids)].add(mx.array(-deltas))
+    if counts_active:
+        ids = np.fromiter(token_counts.keys(), dtype=np.int64, count=len(token_counts))
+        counts = np.fromiter(token_counts.values(), dtype=np.float64, count=len(token_counts))
+        deltas = (frequency * counts + presence * (counts > 0.0)).astype(np.float32)
+        logits = logits.at[mx.array(ids)].add(mx.array(-deltas))
+    if overlay_active:
+        overlay_ids = np.fromiter(
+            penalty_overlay.keys(), dtype=np.int64, count=len(penalty_overlay)
+        )
+        overlay_vals = np.fromiter(
+            penalty_overlay.values(), dtype=np.float64, count=len(penalty_overlay)
+        ).astype(np.float32)
+        logits = logits.at[mx.array(overlay_ids)].add(mx.array(-overlay_vals))
+    return logits
 
 
 class BatchedSparseDistributions:
@@ -85,6 +99,7 @@ def sparse_distribution_from_mlx_logits(
     config: SamplerConfig,
     *,
     token_counts: Mapping[int, int] | None = None,
+    penalty_overlay: Mapping[int, float] | None = None,
 ) -> SparseDistribution | None:
     """Return an exact sparse distribution for top-p then top-k sampling.
 
@@ -96,13 +111,19 @@ def sparse_distribution_from_mlx_logits(
     ``token_counts`` (completion tokens seen so far, scoped by the caller) applies
     the additive presence/frequency penalty to the raw logits BEFORE the
     temperature divide — a no-op (same array) when penalties are 0.
+    ``penalty_overlay`` is the Loop Guard's sparse steering map, applied at the
+    same raw-logit stage.
     """
 
     if config.temperature <= 0 or config.top_k <= 0:
         return None
 
     row = apply_penalties_mlx(
-        logits.reshape(-1), token_counts, config.presence_penalty, config.frequency_penalty
+        logits.reshape(-1),
+        token_counts,
+        config.presence_penalty,
+        config.frequency_penalty,
+        penalty_overlay=penalty_overlay,
     )
     flat = row.astype(mx.float32) / float(config.temperature)
     vocab_size = int(flat.shape[-1])

{mtplx-2.0.0 → mtplx-2.0.2}/mtplx/gdn_capture.py

@@ -1636,9 +1636,13 @@ def gdn_forward_with_capture(
         "yes",
         "on",
     }:
+        from .kernel_selfcheck import lane_disabled
         from .kernels.fused_norm import fused_gdn_norm_gate
 
-        out = fused_gdn_norm_gate(out, z, gdn.norm.weight, gdn.norm.eps)
+        if lane_disabled("fused_gdn_norm_gate"):
+            out = gdn.norm(out, z)
+        else:
+            out = fused_gdn_norm_gate(out, z, gdn.norm.weight, gdn.norm.eps)
     else:
         out = gdn.norm(out, z)
     if not tail_projected:
@@ -1739,15 +1743,20 @@ def forward_with_gdn_capture(
             "yes",
             "on",
         }:
+            from .kernel_selfcheck import lane_disabled
             from .kernels.fused_norm import fused_add_rmsnorm
 
-            h, mlp_input = fused_add_rmsnorm(
-                hidden_states,
-                r,
-                layer.post_attention_layernorm.weight,
-                layer.post_attention_layernorm.eps,
-                threadgroup_size=512,
-            )
+            if lane_disabled("fused_add_rmsnorm"):
+                h = hidden_states + r
+                mlp_input = layer.post_attention_layernorm(h)
+            else:
+                h, mlp_input = fused_add_rmsnorm(
+                    hidden_states,
+                    r,
+                    layer.post_attention_layernorm.weight,
+                    layer.post_attention_layernorm.eps,
+                    threadgroup_size=512,
+                )
         else:
             h = hidden_states + r
             mlp_input = layer.post_attention_layernorm(h)