spb/localvm-research Public License

1# =============================================================================2#  Project   : localvm-research3#  File      : experiments/candidate_01/implementation/runtime.py4#  Purpose   : Margin-gated deferred-refinement generation loop (q4 resident5#              base + windowed q8 verification sweeps with rollback)6#  Author    : Simon-Pierre Boucher7#  Contact   : contact@spboucher.ai8#  Created   : 2026-08-129#  Modified  : 2026-08-1210#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal11#  License   : All rights reserved (research code)12# =============================================================================13"""Candidate-01 runtime.1415Two modes:16  * ``verify-all``  — every window position is checked against the q8 sweep17    (exact q8-greedy output; precision-tiered speculative decoding).18  * ``margin``      — only positions whose base margin < tau are checked19    (approximate; fewer rollbacks — the novel relaxation under test).20"""2122from __future__ import annotations2324import time25from dataclasses import dataclass, field2627import mlx.core as mx28from mlx_lm.models.cache import make_prompt_cache, trim_prompt_cache293031@dataclass32class RunStats:33    tokens_out: int = 034    deferred: int = 035    sweeps: int = 036    rollbacks: int = 037    sweep_time_s: float = 0.038    gen_time_s: float = 0.039    sweep_logical_bytes: int = 040    events: list = field(default_factory=list)414243def _greedy_step(model, inp, cache):44    logits = model(inp, cache=cache)[0, -1].astype(mx.float32)45    top2 = mx.topk(logits, 2)46    nxt = mx.argmax(logits)47    mx.eval(top2, nxt)48    v = top2.tolist()49    return int(nxt.item()), abs(v[1] - v[0])505152def generate_deferred(base_model, verify_model, tokenizer, prompt_ids: list[int],53                      max_tokens: int, tau: float, window: int,54                      mode: str, q8_bytes: int) -> tuple[list[int], RunStats]:55    """Generate with the resident base model, verifying windows with the q856    model. Returns (tokens, stats). `verify_model` weights are assumed57    mmap-backed; each sweep accounts `q8_bytes` logical bytes."""58    assert mode in ("verify-all", "margin")59    stats = RunStats()60    tokens: list[int] = []          # accepted generated tokens61    flags: list[bool] = []          # deferred flag per generated token62    verified_upto = 0               # tokens before this index are settled6364    cache = make_prompt_cache(base_model)65    # Incremental verify cache: sweeps feed only the unseen suffix, cutting66    # sweep compute from O(context) to O(window). v_pos = tokens of the67    # current sequence already ingested; invariant at sweep: v_pos <= base-1.68    # verify_model may be a StreamingVerifier (layer-streamed weights for69    # models that do not fit beside the resident base).70    is_streaming = hasattr(verify_model, "forward_chunk")71    vcache = make_prompt_cache(verify_model.model if is_streaming else verify_model)72    v_pos = 073    t0 = time.perf_counter()74    inp = mx.array(prompt_ids)[None]7576    def sweep() -> bool:77        """Verify tokens[verified_upto:] with q8; rollback on first flip.78        Returns True if a rollback happened."""79        nonlocal verified_upto, tokens, flags, inp, cache, v_pos80        pending = len(tokens) - verified_upto81        if pending <= 0:82            return False83        s0 = time.perf_counter()84        stats.sweeps += 185        stats.sweep_logical_bytes += q8_bytes86        full = prompt_ids + tokens87        base_of_window = len(prompt_ids) + verified_upto88        # feed full[v_pos : end-1]; logits rows are absolute v_pos..end-2,89        # predictions for tokens v_pos+1..end-1 ⊇ the pending window90        chunk = full[v_pos : len(full) - 1]91        if is_streaming:92            logits = verify_model.forward_chunk(chunk, vcache)93        else:94            logits = verify_model(mx.array(chunk)[None], cache=vcache)[0]95        sel = logits[base_of_window - 1 - v_pos :].astype(mx.float32)96        v_pos = len(full) - 197        q8_argmax = mx.argmax(sel, axis=-1)98        mx.eval(q8_argmax)99        q8_next = q8_argmax.tolist()100        flip_at = None101        for j in range(pending):102            i = verified_upto + j103            check = mode == "verify-all" or flags[i]104            if check and q8_next[j] != tokens[i]:105                flip_at = (i, q8_next[j])106                break107        if flip_at is None:108            verified_upto = len(tokens)109            stats.sweep_time_s += time.perf_counter() - s0110            return False111        i, corrected = flip_at112        stats.rollbacks += 1113        # Cache invariant at sweep time: cache = prompt + tokens[:-1] (the114        # newest token was appended but not yet fed). To leave the cache at115        # prompt + tokens[:i], trim (len(tokens)-1 - i) entries.116        trim_prompt_cache(cache, len(tokens) - 1 - i)117        # verify cache must hold a prefix of the corrected sequence:118        # keep exactly prompt + tokens[:i] (= base_next - 1, invariant holds)119        keep = len(prompt_ids) + i120        if v_pos > keep:121            trim_prompt_cache(vcache, v_pos - keep)122            v_pos = keep123        tokens = tokens[:i] + [corrected]124        flags = flags[:i] + [False]125        verified_upto = len(tokens)126        # do NOT feed `corrected` here — the main loop's next _greedy_step127        # feeds tokens[-1], preserving the invariant128        stats.sweep_time_s += time.perf_counter() - s0129        stats.events.append({"rollback_at": i})130        return True131132    # Loop until the full budget is generated AND the tail is verified —133    # a rollback in the final window must resume generation, not truncate.134    max_steps = 6 * max_tokens  # safety bound against rollback ping-pong135    steps = 0136    while steps < max_steps:137        steps += 1138        hit_eos = False139        if len(tokens) < max_tokens:140            nxt, margin = _greedy_step(base_model, inp, cache)141            hit_eos = nxt == tokenizer.eos_token_id142            # EOS is a decision too: append it (always checked) so the sweep143            # can veto a premature stop; stripped before returning.144            tokens.append(nxt)145            flags.append(hit_eos or margin < tau)146            if not hit_eos and margin < tau:147                stats.deferred += 1148            inp = mx.array([[nxt]])149            if not hit_eos and len(tokens) - verified_upto < window and len(tokens) < max_tokens:150                continue151        rolled = sweep()152        if rolled:153            if tokens and tokens[-1] == tokenizer.eos_token_id:154                tokens.pop()  # q8 corrected the decision to "stop here"155                break156            inp = mx.array([[tokens[-1]]])  # corrected token, not yet fed157            continue158        if len(tokens) and tokens[-1] == tokenizer.eos_token_id:159            tokens.pop()160            break161        if len(tokens) >= max_tokens:162            break163164    stats.tokens_out = len(tokens)165    stats.gen_time_s = time.perf_counter() - t0 - stats.sweep_time_s166    return tokens, stats167