Step-2 levers 1+2+3: single-sync all-MLX fused loop

fused_specdecode_generate_mlx — one host sync per block:
- (2) draft ids stay lazy mx tensors and feed the verifier forward
  in-graph (drafter.draft_block_ids + adapter.forward_block_lazy)
- (1) in-graph greedy acceptance (cumprod leading-match) + lazy gather
  of the next-position logits row; per block mx.eval materialises only
  the accept count and candidate ids; drafter-context extensions go
  through mx.async_eval
- (3) no correction forward: the gathered next-row makes the
  verifier's correction the next block's carried bonus, verified (and
  aux-captured) as position 0 of the next batched forward —
  guaranteed-accepted by construction, so every block commits >= 1
  token and the loop can never run below AR pace

Harness uses the new loop automatically on --all-mlx-drafter.

Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

Author: cursoragent

inference_engine/backends/mlx/dflash_drafter.py

@@ -280,31 +280,56 @@ def _run_layers(self, hidden: Any, query_positions: Any, ctx_kv) -> Any:
             hidden = layer(hidden, query_positions, ck, cv)
         return self.norm(hidden)
 
-    def draft_block_cached(
+    def draft_block_ids(
         self,
         ctx_kv,
-        bonus_token_id: int,
+        bonus_id_mx: Any,
         embed_fn: Callable[[Any], Any],
         lm_head_fn: Callable[[Any], Any],
         *,
-        block_size: int,
+        n_masks: int,
         context_len: int,
-    ) -> List[int]:
-        """Single non-causal pass over ``[bonus, mask×block_size]`` against the
-        cached context K/V → ``block_size`` draft token ids. All-MLX: the
-        embed/lm_head fns are the verifier's native MLX weights — no bridge."""
+    ) -> Any:
+        """LAZY draft: ``[bonus, mask×n_masks]`` → mx ``[n_masks]`` draft ids.
+
+        Nothing is evaluated and nothing crosses to python — the returned
+        ids feed the verifier forward inside the same lazy graph (lever ②
+        of the single-sync block loop). ``bonus_id_mx`` is an mx scalar
+        (e.g. ``mx.argmax(next_token_logits)``).
+        """
         mx = _mx()
         cfg = self.cfg
-        query_ids = mx.array(
-            [[int(bonus_token_id)] + [cfg.mask_token_id] * block_size])
-        query_positions = mx.arange(context_len, context_len + 1 + block_size)
+        mask_ids = mx.full((n_masks,), cfg.mask_token_id, dtype=bonus_id_mx.dtype)
+        query_ids = mx.concatenate([bonus_id_mx[None], mask_ids])[None]
+        query_positions = mx.arange(context_len, context_len + 1 + n_masks)
         h = embed_fn(query_ids).astype(self.fc.weight.dtype)
         h = self._run_layers(h, query_positions, ctx_kv)
-        logits = lm_head_fn(h)  # [1, 1+block, vocab]
+        logits = lm_head_fn(h)  # [1, 1+n_masks, vocab]
         vocab = logits.shape[-1]
         never_mask = mx.arange(vocab) == cfg.mask_token_id
         logits = mx.where(never_mask, mx.array(-float("inf")), logits)
-        drafts = mx.argmax(logits[0, 1:1 + block_size], axis=-1)
+        return mx.argmax(logits[0, 1:1 + n_masks], axis=-1)
+
+    def draft_block_cached(
+        self,
+        ctx_kv,
+        bonus_token_id: int,
+        embed_fn: Callable[[Any], Any],
+        lm_head_fn: Callable[[Any], Any],
+        *,
+        block_size: int,
+        context_len: int,
+    ) -> List[int]:
+        """Single non-causal pass over ``[bonus, mask×block_size]`` against the
+        cached context K/V → ``block_size`` draft token ids (materialised).
+        Compatibility surface for the generic fused loop / parity gate; the
+        single-sync loop uses :meth:`draft_block_ids` instead."""
+        mx = _mx()
+        drafts = self.draft_block_ids(
+            ctx_kv, mx.array(int(bonus_token_id), dtype=mx.uint32),
+            embed_fn, lm_head_fn,
+            n_masks=block_size, context_len=context_len,
+        )
         mx.eval(drafts)
         return [int(t) for t in drafts.tolist()]
 

inference_engine/backends/mlx/fused_specdecode.py

@@ -223,6 +223,33 @@ def last_aux_torch_slice(self, start: int = 0, end: Optional[int] = None) -> Lis
         bridge = self._bridge or (lambda a: a)
         return [bridge(a[start:end]) for a in self._last_aux_mx]
 
+    def forward_block_lazy(self, ids_mx: Any) -> Any:
+        """LAZY incremental verify: ``ids_mx`` is an mx ``[1, L]`` (typically
+        the in-graph concatenation of the carried bonus + lazy draft ids —
+        lever ② of the single-sync loop). Returns ``mx [L, V]`` logits with
+        NO evaluation; aux hidden (when ``_capture_aux``) stays lazy in
+        ``_last_aux_mx`` and is consumed lazily by the drafter-context
+        extension."""
+        if self._cache is None:
+            raise RuntimeError("verifier not prefilled")
+        want_aux = self._capture_aux and bool(self.aux_layer_ids)
+        if want_aux:
+            sink: Dict[int, Any] = {}
+            with _patched_decoder_layers(self.text_model):
+                for layer in self.text_model.layers:
+                    layer._kakeya_aux_sink = sink
+                    layer._aux_record = sink
+                logits = self.mlx_model(ids_mx, cache=self._cache)
+                aux = _build_aux(self.text_model, ids_mx, sink,
+                                 self.embed_scale, self.aux_layer_ids)
+            self._last_aux_mx = [a[0] for a in aux]  # [L, hidden] each, lazy
+            self._last_aux = None
+        else:
+            logits = self.mlx_model(ids_mx, cache=self._cache)
+            self._last_aux = None
+            self._last_aux_mx = None
+        return logits[0]
+
     def commit_or_truncate(self, *, forwarded: int, accepted: int) -> None:
         if accepted < 0 or accepted > forwarded:
             raise ValueError("accepted must satisfy 0 <= accepted <= forwarded")
@@ -277,6 +304,125 @@ def lm_head_fn(h: Any) -> Any:
     return embed_fn, lm_head_fn
 
 
+# --------------------------------------------------------------------------- #
+# Single-sync all-MLX fused loop (levers ① ② ③ of the Step-2 throughput plan;
+# docs/mlx-port-lessons.md "Step-2 rescue status").
+# --------------------------------------------------------------------------- #
+def fused_specdecode_generate_mlx(
+    adapter: "MLXRestoredIncrementalVerifier",
+    drafter: Any,
+    *,
+    aux_prompt: Sequence[Any],
+    embed_fn: Callable[[Any], Any],
+    lm_head_fn: Callable[[Any], Any],
+    gen_tokens: int,
+    block_size: int,
+    eos_ids: Sequence[int] = (),
+) -> Dict[str, Any]:
+    """All-MLX fused spec decode with ONE host sync per block.
+
+    * ② draft+verify single graph: the drafter's lazy draft ids
+      (:meth:`MLXDFlashDrafter.draft_block_ids`) are concatenated with the
+      carried bonus in-graph and fed straight into the verifier forward —
+      no draft token ever crosses to python before verification.
+    * ① in-graph acceptance: the leading-match count is
+      ``sum(cumprod(argmax(pred_rows) == candidate))``; the next-position
+      logits row is gathered with the lazy count (``mx.take``). The block's
+      single ``mx.eval`` materialises exactly three things: the accept
+      count, the candidate ids, and nothing else. Drafter-context
+      extensions are pushed with ``mx.async_eval`` so Metal works while
+      python does bookkeeping.
+    * ③ carried correction: on rejection there is NO correction forward.
+      ``next_token_logits`` is set to the gathered next-position row, so
+      the verifier's own argmax (the correction) becomes the next block's
+      bonus — guaranteed-accepted at position 0 of the next verify, where
+      its K/V and aux are computed as part of the batched forward.
+
+    Per-block commit = the accepted candidate prefix (position 0, the
+    carried bonus/correction, always accepts by construction — every block
+    commits >= 1 token, so the loop degrades to AR pace, never below).
+    """
+    import mlx.core as mx  # type: ignore
+
+    eos = set(int(t) for t in eos_ids)
+    C = adapter._past_len
+    t_ctx = time.perf_counter()
+    ctx_kv = drafter.make_context_kv(list(aux_prompt), mx.arange(0, C))
+    mx.async_eval([t for kv in ctx_kv for t in kv])
+    timing = {
+        "ctx_kv_build_s": time.perf_counter() - t_ctx,
+        "build_s": 0.0,   # lazy graph construction (python-side)
+        "eval_s": 0.0,    # the per-block single sync (Metal compute)
+        "extend_s": 0.0,
+    }
+    adapter._capture_aux = True
+
+    generated: List[int] = []
+    accepts: List[int] = []
+    next_logits = adapter.next_token_logits  # mx [V], may be lazy
+    try:
+        while len(generated) < gen_tokens:
+            L = min(block_size, gen_tokens - len(generated))
+            cstart = adapter._past_len
+            t_build = time.perf_counter()
+            bonus_id = mx.argmax(next_logits)            # lazy scalar
+            n_draft = max(L - 1, 0)
+            if n_draft:
+                drafts = drafter.draft_block_ids(
+                    ctx_kv, bonus_id, embed_fn, lm_head_fn,
+                    n_masks=n_draft, context_len=cstart)
+                candidate = mx.concatenate([bonus_id[None], drafts])  # [L]
+            else:
+                candidate = bonus_id[None]
+            block_logits = adapter.forward_block_lazy(candidate[None])  # [L, V]
+            # In-graph greedy acceptance: row i of pred_rows predicts
+            # candidate[i]; leading-match count via cumprod.
+            pred_rows = mx.concatenate(
+                [next_logits[None], block_logits[:-1]], axis=0)  # [L, V]
+            matches = (mx.argmax(pred_rows, axis=-1) == candidate)
+            accepted_mx = mx.sum(mx.cumprod(matches.astype(mx.int32)))
+            # Logits predicting position cstart+accepted (the carried
+            # bonus/correction source for the next block).
+            rows = mx.concatenate([next_logits[None], block_logits], axis=0)
+            next_row = mx.take(rows, accepted_mx[None], axis=0)[0]   # [V]
+            timing["build_s"] += time.perf_counter() - t_build
+            # ---- the block's single host sync ----
+            t_eval = time.perf_counter()
+            mx.eval(accepted_mx, candidate)
+            timing["eval_s"] += time.perf_counter() - t_eval
+            accepted = int(accepted_mx.item())
+            cand = [int(x) for x in candidate.tolist()]
+            adapter.commit_or_truncate(forwarded=L, accepted=accepted)
+            commit = cand[:accepted]
+            generated += commit
+            accepts.append(accepted)
+            next_logits = next_row
+            adapter.next_token_logits = next_row
+            if accepted and adapter._last_aux_mx is not None:
+                t_extend = time.perf_counter()
+                new_aux = [a[0:accepted][None] for a in adapter._last_aux_mx]
+                ctx_kv = drafter.extend_context_kv(
+                    ctx_kv,
+                    drafter.make_context_kv(
+                        new_aux, mx.arange(cstart, cstart + accepted)))
+                mx.async_eval([t for kv in ctx_kv for t in kv])
+                timing["extend_s"] += time.perf_counter() - t_extend
+            if any(t in eos for t in commit):
+                break
+    finally:
+        adapter._capture_aux = False
+    generated = generated[:gen_tokens]
+    return {
+        "tokens": generated,
+        "blocks": len(accepts),
+        "mean_accept_len": (round(sum(accepts) / len(accepts), 3)
+                            if accepts else 0.0),
+        "decode_tokens": len(generated),
+        "loop": "mlx_single_sync_v2",
+        "time_breakdown_s": {k: round(v, 3) for k, v in timing.items()},
+    }
+
+
 # --------------------------------------------------------------------------- #
 # The fused spec-decode loop (control flow; MLX/torch ops via injected fns).
 # --------------------------------------------------------------------------- #

scripts/research/k3_integrated_niah_eval_mac.py

@@ -179,6 +179,7 @@ def main() -> int:
     from inference_engine.backends.mlx.fused_specdecode import (
         MLXRestoredIncrementalVerifier, capture_aux_hidden,
         make_bridge_embed_lm_head, fused_specdecode_generate,
+        fused_specdecode_generate_mlx,
     )
     from inference_engine.v04.kv_compressor import make_default_compressor
     from inference_engine.bench.k3_report_gate import (
@@ -684,13 +685,21 @@ def eval_fused_specdecode() -> Tuple[List[str], List[float], List[int]]:
             prefill_s = time.perf_counter() - prefill_t0
             t0 = time.perf_counter()
             if args.force_fused_specdecode:
-                res = fused_specdecode_generate(
-                    adapter, active_drafter, aux_prompt=aux_prompt,
-                    embed_fn=embed_fn, lm_head_fn=lm_head_fn,
-                    gen_tokens=args.max_new_tokens, block_size=args.block_size,
-                    eos_ids=end_ids,
-                    argmax_fn=argmax_fn, arange_fn=arange_fn, cat_aux_fn=cat_aux_fn,
-                    allow_greedy_fallback=False)
+                if mlx_drafter is not None:
+                    # Single-sync all-MLX loop (levers ①②③).
+                    res = fused_specdecode_generate_mlx(
+                        adapter, active_drafter, aux_prompt=aux_prompt,
+                        embed_fn=embed_fn, lm_head_fn=lm_head_fn,
+                        gen_tokens=args.max_new_tokens,
+                        block_size=args.block_size, eos_ids=end_ids)
+                else:
+                    res = fused_specdecode_generate(
+                        adapter, active_drafter, aux_prompt=aux_prompt,
+                        embed_fn=embed_fn, lm_head_fn=lm_head_fn,
+                        gen_tokens=args.max_new_tokens, block_size=args.block_size,
+                        eos_ids=end_ids,
+                        argmax_fn=argmax_fn, arange_fn=arange_fn, cat_aux_fn=cat_aux_fn,
+                        allow_greedy_fallback=False)
                 res["drafter_runtime"] = "mlx" if mlx_drafter is not None else "torch"
             else:
                 t_greedy = time.perf_counter()