evidence/source/e3_rollout_v3/sample.py

"""
Generate from a trained checkpoint, and measure MTP draft acceptance.

Two modes:

  --mode ar    plain autoregressive decoding with a KV cache. This is your
               throughput baseline.

  --mode mtp   self-speculative decoding: the MTP module drafts K tokens, the
               trunk verifies them in one batched forward, and acceptance uses
               the Leviathan-Chen probability ratio with residual correction
               (p - q)+ on rejection. Reported per-depth acceptance is the number
               that tells you whether the heads actually trained.

The mtp path reuses a successful verification forward as the next cycle's target
state. After a rejection it recomputes the corrected full prefix rather than
rolling back a KV cache. That makes it a correctness and rollout-acceptance
reference, not the final speed path. Speed is MTPLX's job once the model is
exported.

Usage:
    python sample.py --ckpt out/run1/ckpt_latest --mode mtp --depth 3 \\
        --prompt "def quicksort(arr):"
"""

import argparse
import json
import os
import time

import mlx.core as mx
import numpy as np
from mlx.utils import tree_map, tree_unflatten
from tokenizers import Tokenizer

from model import Wisp, ModelArgs, causal_mask
from scripts.rollout_metrics import normalized_entropy, token_ids_sha256

DTYPES = {"float32": mx.float32, "bfloat16": mx.bfloat16, "float16": mx.float16}


def load_model(ckpt_dir: str, compute_dtype):
    with open(os.path.join(ckpt_dir, "meta.json")) as f:
        meta = json.load(f)
    args = ModelArgs.from_dict(meta["model_args"])
    model = Wisp(args)
    master = tree_unflatten(list(mx.load(os.path.join(ckpt_dir, "master.safetensors")).items()))
    model.update(tree_map(
        lambda a: a.astype(compute_dtype) if isinstance(a, mx.array) else a, master
    ))
    model.eval()
    mx.eval(model.parameters())
    return model, args, meta


def to_probs(logits_row: np.ndarray, temperature: float, top_k: int, top_p: float) -> np.ndarray:
    """Turn one row of logits into the exact distribution we will sample from."""
    x = logits_row.astype(np.float64)
    if temperature <= 0:
        probs = np.zeros_like(x)
        probs[int(np.argmax(x))] = 1.0
        return probs

    x = x / temperature
    if top_k and top_k < x.shape[0]:
        cutoff = np.partition(x, -top_k)[-top_k]
        x = np.where(x < cutoff, -np.inf, x)

    x = x - np.max(x)
    probs = np.exp(x)
    probs /= probs.sum()

    if top_p and top_p < 1.0:
        order = np.argsort(-probs)
        cumulative = np.cumsum(probs[order])
        # Conventional nucleus sampling keeps the smallest set whose mass is at
        # least top_p, which means the token that *crosses* the threshold is
        # included. `cumulative <= top_p` drops it, so the retained mass is below
        # top_p and "top_p=0.95" silently means something tighter than 0.95.
        # Acceptance stays internally exact either way, but the knob has to mean
        # what everyone else means by it.
        crossing = int(np.searchsorted(cumulative, top_p))
        keep = np.zeros_like(cumulative, dtype=bool)
        keep[:min(crossing + 1, keep.size)] = True
        mask = np.zeros_like(probs, dtype=bool)
        mask[order[keep]] = True
        probs = np.where(mask, probs, 0.0)
        probs /= probs.sum()

    return probs


def sample_from(probs: np.ndarray, rng: np.random.Generator) -> int:
    return int(rng.choice(probs.shape[0], p=probs))


def generate_ar(model, args, prompt_ids, max_tokens, temperature, top_k, top_p, rng):
    tokens = list(prompt_ids)
    caches = None
    fed = mx.array([tokens], dtype=mx.int32)
    dtype = model.norm.weight.dtype

    start = time.time()
    generated = 0
    for _ in range(max_tokens):
        length = fed.shape[1]
        mask = causal_mask(length, dtype) if length > 1 else None
        if caches is not None and length == 1:
            mask = None
        logits, _, caches = model(fed, mask, caches)
        mx.eval(logits)
        probs = to_probs(np.array(logits[0, -1].astype(mx.float32), copy=False), temperature, top_k, top_p)
        nxt = sample_from(probs, rng)
        tokens.append(nxt)
        generated += 1
        fed = mx.array([[nxt]], dtype=mx.int32)
    elapsed = time.time() - start
    return tokens, generated / max(elapsed, 1e-9)


def generate_mtp(model, args, prompt_ids, max_tokens, depth, temperature, top_k, top_p,
                 draft_temperature, rng, policy="fixed",
                 entropy_threshold=0.5, capture_trace=False):
    """
    Self-speculative decoding with either a fixed draft depth or an adaptive one.

    The adaptive policy exists because draft depth is a bet, not a constant. Every
    drafted token costs a recursion of the MTP module and a column of the verify
    forward, and a draft sampled from a flat distribution is very likely to be
    rejected, which throws away that cost and everything drafted behind it. When
    the drafter is confident, deeper drafting is nearly free yield. When it is not,
    the marginal draft is negative expected value.

    So the policy drafts one token at a time and keeps going only while the
    drafter's normalized full-vocabulary entropy stays at or below
    `entropy_threshold`, up to `depth`. The first draft is always issued. The
    signal needs no extra model call because its logits are already available.

    A successful verification already contains both the hidden state and target
    logits needed to start the next cycle. Those are reused when every draft is
    accepted. A rejection replaces one draft with a residual-correction token, so
    hidden states from that position onward are invalid; the next cycle recomputes
    the corrected full prefix. The target-forward count is therefore one initial
    prefix pass, one pass per verification, and one recovery pass after each
    rejection that is followed by more generation.

    This still has no rollback-capable KV cache, so wall-clock throughput is a
    reference rather than the deployment result. The reported
    `mean_accepted_per_verification` and rollout acceptance by depth are real
    sequential histories, unlike the teacher-forced E1 diagnostic.

    Set `capture_trace=True` to include a strict cycle-level account of the
    generation in `stats["generation_trace"]`. Trace hashes bind each state to
    the exact token IDs consumed by the corresponding model forward.
    """
    tokens = list(prompt_ids)
    prompt_token_count = len(tokens)
    dtype = model.norm.weight.dtype
    accepted_at_depth = np.zeros(depth, dtype=np.int64)
    trials_at_depth = np.zeros(depth, dtype=np.int64)
    drafts_issued = 0
    draft_recursions = 0
    corrections = 0
    cycles = 0
    prefix_forwards = 0
    recovery_forwards = 0
    verification_forwards = 0
    fully_accepted_verifications = 0
    draft_entropies = [[] for _ in range(depth)]
    prefix_hidden = None
    next_target_logits = None
    generation_cycles = [] if capture_trace else None

    start = time.time()
    produced = 0
    while produced < max_tokens:
        cycle_index = cycles
        output_start = produced
        prefix_token_count = len(tokens)
        state_source = (
            "full_recompute"
            if prefix_hidden is None
            else "verification_reuse"
        )
        cycles += 1
        if prefix_hidden is None:
            seq = mx.array([tokens], dtype=mx.int32)
            mask = causal_mask(seq.shape[1], dtype)
            prefix_hidden, _ = model.trunk(seq, mask)
            logits = model.head(prefix_hidden[:, -1:, :])
            mx.eval(prefix_hidden, logits)
            next_target_logits = np.array(
                logits[0, -1].astype(mx.float32), copy=False
            )
            prefix_forwards += 1
            if cycles > 1:
                recovery_forwards += 1
        hidden = prefix_hidden

        # Bonus token: comes straight from the target, always valid.
        p0 = to_probs(
            next_target_logits, temperature, top_k, top_p
        )
        bonus = sample_from(p0, rng)
        tokens.append(bonus)
        produced += 1
        emitted_token_ids = [bonus]
        if produced >= max_tokens:
            if capture_trace:
                generation_cycles.append({
                    "cycle_index": cycle_index,
                    "output_start": output_start,
                    "prefix_token_count": prefix_token_count,
                    "prefix_sha256": token_ids_sha256(
                        int(token) for token in tokens[:-1]
                    ),
                    "state_source": state_source,
                    "bonus_token": bonus,
                    "draft_attempts": [],
                    "verification": None,
                    "emitted_token_ids": emitted_token_ids,
                    "next_state": "terminal",
                })
            break

        # Draft `depth` tokens by recursing the MTP module over the whole prefix.
        #
        # The module contains a transformer block, and during training that block
        # attends across the full window under a causal mask. Handing it a single
        # position at inference would leave its attention with nothing but itself
        # to look at, which is a different function from the one that was trained
        # and shows up as depressed acceptance rather than as an error. Recomputing
        # over the prefix is the same tradeoff this mode already makes for the
        # trunk: correctness reference first, speed is MTPLX's job.
        drafts, draft_probs = [], []
        draft_attempts = []
        prefix_len = hidden.shape[1]
        mtp_mask = causal_mask(prefix_len, dtype)
        cur = hidden
        # Position i consumes the embedding of token i+k at depth k, mirroring
        # model.loss. At entry that window is token indices 1 through prefix_len.
        window = tokens[1:prefix_len + 1]
        for _ in range(depth):
            emb = model.tok_emb(mx.array([window], dtype=mx.int32))
            cur, _ = model.mtp(cur, emb, mtp_mask)
            dl = model.head(cur[:, -1:, :])
            mx.eval(dl)
            draft_recursions += 1
            draft_logits = np.array(
                dl[0, -1].astype(mx.float32), copy=False
            )
            policy_probs = to_probs(draft_logits, 1.0, 0, 1.0)
            entropy = normalized_entropy(policy_probs)
            draft_index = len(drafts)
            draft_entropies[draft_index].append(entropy)
            if (
                policy == "adaptive"
                and draft_index > 0
                and entropy > entropy_threshold
            ):
                if capture_trace:
                    draft_attempts.append({
                        "depth": draft_index + 1,
                        "normalized_entropy": float(entropy),
                        "issued": False,
                        "draft_token": None,
                    })
                break
            q = to_probs(
                draft_logits,
                draft_temperature,
                top_k,
                top_p,
            )
            tok = sample_from(q, rng)
            drafts.append(tok)
            draft_probs.append(q)
            if capture_trace:
                draft_attempts.append({
                    "depth": draft_index + 1,
                    "normalized_entropy": float(entropy),
                    "issued": True,
                    "draft_token": tok,
                })
            window = window[1:] + [tok]

        # Verify all drafts in one batched trunk forward.
        candidate = tokens + drafts
        seq = mx.array([candidate], dtype=mx.int32)
        mask = causal_mask(seq.shape[1], dtype)
        vh, _ = model.trunk(seq, mask)
        base = len(tokens) - 1  # position whose logits predict drafts[0]
        # The final row predicts the next cycle's bonus when every draft is
        # accepted, so compute it in the same projection.
        vlogits = model.head(vh[:, base:base + len(drafts) + 1, :])
        mx.eval(vh, vlogits)
        vlogits_np = np.array(vlogits[0].astype(mx.float32), copy=False)
        verification_forwards += 1
        drafts_issued += len(drafts)

        accepted_this_cycle = 0
        rejection_depth = None
        outcomes = []
        for j in range(len(drafts)):
            trials_at_depth[j] += 1
            p = to_probs(vlogits_np[j], temperature, top_k, top_p)
            q = draft_probs[j]
            tok = drafts[j]
            ratio = 1.0 if q[tok] <= 0 else min(1.0, p[tok] / q[tok])
            if rng.random() < ratio:
                accepted_at_depth[j] += 1
                accepted_this_cycle += 1
                tokens.append(tok)
                produced += 1
                emitted_token = tok
                accepted = True
                if produced >= max_tokens:
                    if capture_trace:
                        emitted_token_ids.append(emitted_token)
                        outcomes.append({
                            "depth": j + 1,
                            "draft_token": tok,
                            "target_argmax_token": int(
                                np.argmax(vlogits_np[j])
                            ),
                            "target_row_index": base + j,
                            "accepted": accepted,
                            "emitted_token": emitted_token,
                        })
                    break
            else:
                residual = np.maximum(p - q, 0.0)
                total = residual.sum()
                residual = p if total <= 0 else residual / total
                emitted_token = sample_from(residual, rng)
                tokens.append(emitted_token)
                corrections += 1
                produced += 1
                accepted = False
                rejection_depth = j + 1
            if capture_trace:
                emitted_token_ids.append(emitted_token)
                outcomes.append({
                    "depth": j + 1,
                    "draft_token": tok,
                    "target_argmax_token": int(np.argmax(vlogits_np[j])),
                    "target_row_index": base + j,
                    "accepted": accepted,
                    "emitted_token": emitted_token,
                })
            if not accepted:
                break

        fully_accepted = accepted_this_cycle == len(drafts)
        if fully_accepted:
            fully_accepted_verifications += 1
            if produced < max_tokens:
                prefix_hidden = vh
                next_target_logits = vlogits_np[-1]
        elif produced < max_tokens:
            # The correction token was not part of `vh`; recompute its prefix at
            # the start of the next cycle.
            prefix_hidden = None
            next_target_logits = None

        if capture_trace:
            generation_cycles.append({
                "cycle_index": cycle_index,
                "output_start": output_start,
                "prefix_token_count": prefix_token_count,
                "prefix_sha256": token_ids_sha256(
                    int(token)
                    for token in candidate[:prefix_token_count]
                ),
                "state_source": state_source,
                "bonus_token": bonus,
                "draft_attempts": draft_attempts,
                "verification": {
                    "candidate_sha256": token_ids_sha256(
                        int(token) for token in candidate
                    ),
                    "base_row_index": base,
                    "projected_row_count": len(drafts) + 1,
                    "outcomes": outcomes,
                    "rejection_depth": rejection_depth,
                    "fully_accepted": fully_accepted,
                    "unused_drafts": len(drafts) - len(outcomes),
                },
                "emitted_token_ids": emitted_token_ids,
                "next_state": (
                    "terminal"
                    if produced >= max_tokens
                    else (
                        "verification_reuse"
                        if fully_accepted
                        else "full_recompute"
                    )
                ),
            })

    elapsed = time.time() - start
    rollout_acceptance = [
        round(float(a) / float(t), 4) if t else None
        for a, t in zip(accepted_at_depth, trials_at_depth)
    ]
    target_forwards = prefix_forwards + verification_forwards
    accepted_drafts = int(accepted_at_depth.sum())
    stats = {
        "policy": policy,
        "entropy_threshold": (
            entropy_threshold if policy == "adaptive" else None
        ),
        "cycles": cycles,
        "tokens": produced,
        "corrections": corrections,
        "drafts_issued": int(drafts_issued),
        "draft_recursions": int(draft_recursions),
        "elapsed_seconds": elapsed,
        "tok_per_sec": produced / max(elapsed, 1e-9),
        "acceptance_per_depth": rollout_acceptance,
        "rollout_acceptance_per_depth": rollout_acceptance,
        "rollout_accepted_per_depth": [
            int(value) for value in accepted_at_depth
        ],
        "rollout_trials_per_depth": [
            int(value) for value in trials_at_depth
        ],
        "prefix_forwards": prefix_forwards,
        "recovery_forwards": recovery_forwards,
        "verification_forwards": verification_forwards,
        "target_forwards": target_forwards,
        "fully_accepted_verifications": fully_accepted_verifications,
        "mean_accepted_per_verification": round(
            accepted_drafts / max(verification_forwards, 1), 3
        ),
        "mean_output_tokens_per_target_forward": round(
            produced / max(target_forwards, 1), 3
        ),
        "mean_accepted_per_cycle": round(
            accepted_drafts / max(cycles, 1), 3
        ),
        "mean_tokens_per_cycle": round(produced / max(cycles, 1), 3),
        "mean_drafts_per_cycle": round(drafts_issued / max(cycles, 1), 3),
        "mean_draft_recursions_per_cycle": round(
            draft_recursions / max(cycles, 1), 3
        ),
        "draft_yield": round(
            accepted_drafts / max(drafts_issued, 1), 4
        ),
        "mean_normalized_entropy_per_depth": [
            (
                round(float(np.mean(values)), 6)
                if values else None
            )
            for values in draft_entropies
        ],
    }
    if capture_trace:
        stats["generation_trace"] = {
            "schema_version": 1,
            "prompt_token_count": prompt_token_count,
            "max_tokens": max_tokens,
            "cycles": generation_cycles,
        }
    return tokens, stats


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--tokenizer", default="tokenizer/code32k.json")
    ap.add_argument("--prompt", default="def quicksort(arr):")
    ap.add_argument("--mode", choices=["ar", "mtp"], default="ar")
    ap.add_argument("--max-tokens", type=int, default=128)
    ap.add_argument("--depth", type=int, default=2,
                    help="fixed draft depth, or the cap under --policy adaptive")
    ap.add_argument("--policy", choices=["fixed", "adaptive"], default="fixed",
                    help="adaptive stops drafting when normalized entropy is high")
    ap.add_argument(
        "--entropy-threshold",
        type=float,
        default=0.5,
        help="adaptive policy: draft deeper while normalized H(q) stays below this",
    )
    ap.add_argument("--temperature", type=float, default=0.6)
    ap.add_argument("--draft-temperature", type=float, default=0.7)
    ap.add_argument("--top-k", type=int, default=20)
    ap.add_argument("--top-p", type=float, default=0.95)
    ap.add_argument("--dtype", default="bfloat16")
    ap.add_argument("--seed", type=int, default=0)
    cli = ap.parse_args()
    if cli.depth < 1:
        raise ValueError("--depth must be positive")
    if not 0.0 <= cli.entropy_threshold <= 1.0:
        raise ValueError("--entropy-threshold must be between 0 and 1")

    rng = np.random.default_rng(cli.seed)
    compute_dtype = DTYPES[cli.dtype]
    model, args, meta = load_model(cli.ckpt, compute_dtype)
    tok = Tokenizer.from_file(cli.tokenizer)
    prompt_ids = tok.encode(cli.prompt).ids

    print(f"checkpoint step {meta['step']}, {args.n_layers}L/{args.dim}d, "
          f"mtp_layers={args.mtp_layers}")

    if cli.mode == "ar":
        tokens, tps = generate_ar(
            model, args, prompt_ids, cli.max_tokens,
            cli.temperature, cli.top_k, cli.top_p, rng,
        )
        print(f"\n{tok.decode(tokens)}\n")
        print(json.dumps({"mode": "ar", "tok_per_sec": round(tps, 2)}))
    else:
        tokens, stats = generate_mtp(
            model, args, prompt_ids, cli.max_tokens, cli.depth,
            cli.temperature, cli.top_k, cli.top_p, cli.draft_temperature, rng,
            policy=cli.policy, entropy_threshold=cli.entropy_threshold,
        )
        print(f"\n{tok.decode(tokens)}\n")
        print(json.dumps({"mode": "mtp", "depth": cli.depth, **stats}))


if __name__ == "__main__":
    main()