""" 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()