Retire the Qwen AR pipeline

tools/issue104-reference.py

@@ -0,0 +1,66 @@
+"""Pinned MTPLX full-model classic AR receipts. Reference only; no downloads.
+
+Run serially under test-supervisor with MLX 0.32.2 and pinned MTPLX on PYTHONPATH.
+"""
+import hashlib
+import json
+from pathlib import Path
+import sys
+
+print("#104 importing pinned MTPLX", file=sys.stderr, flush=True)
+import mlx.core as mx
+import numpy as np
+from mlx_lm.utils import load_model
+from mtplx.models import qwen4_exp
+from mtplx.runtime import MTPLXRuntime, _load_tokenizer_resilient
+from mtplx.sampling import SamplerConfig
+from mtplx.generation import generate_ar
+from mtplx.cache_state import snapshot_cache
+from mtplx.session_bank import SessionBank
+
+assert mx.__version__ == "0.32.2"
+root = Path(sys.argv[1])
+print("#104 loading unchanged Bare-Speed model", file=sys.stderr, flush=True)
+model, config = load_model(root, lazy=False, strict=True,
+    get_model_classes=lambda **_: (qwen4_exp.Model, qwen4_exp.ModelArgs))
+model.post_weight_load(root)
+tokenizer = _load_tokenizer_resilient(root, config)
+rt = MTPLXRuntime(model, tokenizer, root, False, None)
+
+def fingerprint(array):
+    values = np.asarray(array.astype(mx.float32)).astype('<f4')
+    assert np.isfinite(values).all()
+    return dict(shape=list(array.shape), sha256=hashlib.sha256(values.tobytes()).hexdigest())
+
+def final_state(out, prompt):
+    assert out.final_state is not None and out.final_state.safe_to_commit
+    final = out.final_state
+    snapshot = snapshot_cache(final.final_trunk_cache)
+    states = [[None if a is None else fingerprint(a) for a in arrays]
+              if arrays is not None else None for arrays in snapshot.states]
+    return dict(ids=prompt + out.tokens, logits=fingerprint(final.final_logits), states=states)
+
+records, warm_records = [], []
+for temperature in (.6, 1.):
+    prompt = [9707, 11, 1075]
+    sampler = SamplerConfig(temperature=temperature, top_p=.95, top_k=20)
+    bank = SessionBank()
+    for warm in (False, True):
+        out = generate_ar(rt, prompt, max_tokens=8,
+            sampler=sampler, seed=0, stop_token_ids=set(), capture_final_state=True,
+            session_bank=bank, session_id="issue104", store_prefix_snapshot=False,
+            token_callback=lambda ids: print(f"#104 t={temperature} warm={warm}: {ids}", file=sys.stderr, flush=True))
+        value = final_state(out, prompt)
+        (warm_records if warm else records).append(dict(temperature=temperature, state=value))
+        print(json.dumps(dict(event="issue104_oracle", temperature=temperature, warm=warm,
+            tokens=out.tokens, state=value, cached_tokens=out.stats.cached_tokens,
+            seconds=out.stats.elapsed_s)), flush=True)
+        if not warm:
+            final = out.final_state
+            assert bank.put(runtime=rt, token_ids=value["ids"], cache=final.final_trunk_cache,
+                logits=final.final_logits, hidden=None, hidden_variant="post_norm",
+                session_id="issue104", mtp_history_policy="cycle", snapshot_epoch=len(value["ids"])) is not None
+            prompt = value["ids"] + [1075]
+        else:
+            assert out.stats.cached_tokens == len(prompt) - 1
+print(json.dumps(dict(event="issue104_classic_baseline", records=records, warm_records=warm_records)), flush=True)

tools/mtplx-binding-reference.py

@@ -48,7 +48,7 @@ def main(root, mode=None):
                     hot_mb=sidecar._hot_cap_rows * sidecar._hot_row_bytes // 2**20,
                     prefetch=sidecar._pool is not None, shape=list(out.shape),
                     sha256=hashlib.sha256(np.asarray(out.astype(mx.float32)).astype('<f4').tobytes()).hexdigest()))
-            print(json.dumps(dict(mask=mask, mtp=mtp, tables=tables, ar_ready=model.set_ar_pipeline_mode(True)),separators=(",", ":")),flush=True)
+            print(json.dumps(dict(mask=mask, mtp=mtp, tables=tables),separators=(",", ":")),flush=True)
             del model
             continue
         parameters = tree_flatten(model.parameters())

tools/mtplx-kernel-fixtures.py

@@ -1627,7 +1627,7 @@ def ar_driver_fixtures():
     """Unchanged generate_ar with a tiny GPU transition-table model, no weights.
 
     Isolates orchestration (not model or throughput parity); actual prefill,
-    samplers, pipeline and final-state capture are all MTPLX implementations.
+    samplers and final-state capture are all MTPLX implementations.
     """
     from mtplx.generation import generate_ar
     from mtplx.runtime import MTPLXRuntime
@@ -1641,11 +1641,8 @@ def ar_driver_fixtures():
     table = mx.array(values.reshape(vocab, vocab))
     class Model:
         def __init__(self):
-            self.calls, self.modes = [], []
+            self.calls = []
         def make_cache(self): return []
-        def set_ar_pipeline_mode(self, enabled):
-            self.modes.append(bool(enabled))
-            return True
         def __call__(self, ids, cache=None, return_hidden=False, emit_logits=True, logits_keep=None):
             self.calls.append((ids, return_hidden))
             out = table[ids]
@@ -1653,12 +1650,11 @@ def ar_driver_fixtures():
     tokenizer = SimpleNamespace(decode=lambda ids, **kwargs: " ".join(map(str, ids)))
     for pipeline in (False, True):
         for asynchronous in (False, True):
-            for temperature, penalty in ((0., 0.), (.6, 0.), (.6, .4)):
+            for temperature, top_p, top_k, penalty in ((0., .95, 20, 0.), (.6, .95, 20, 0.), (.6, .95, 20, .4), (1., .95, 20, 0.), (1., .5, 2, 0.), (1., 0., 2, 0.), (1., 1., 2, 0.)):
                 for limit in (0, 1, 2, 3, 12):
                     for capture in (False, True):
-                        # 13/35 are the first/second sampled pipeline tokens:
-                        # exercise immediate EOS and EOS with an in-flight next
-                        # graph, not just limits and a non-triggering stop set.
+                        # Retain the historical stop IDs and explicitly set the
+                        # retired flag to prove current MTPLX ignores it.
                         for stop_ids in (set(), {0, 5, 36}, {13}, {35}):
                             os.environ["MTPLX_AR_PIPELINE"] = "1" if pipeline else "0"
                             os.environ["MTPLX_ASYNC_AR"] = "1" if asynchronous else "0"
@@ -1666,15 +1662,15 @@ def ar_driver_fixtures():
                             rt = MTPLXRuntime(model, tokenizer, Path("."), False, None)
                             emitted = []
                             out = generate_ar(rt, [1, 2], max_tokens=limit,
-                                sampler=SamplerConfig(temperature=temperature, top_p=.95, top_k=20, presence_penalty=penalty),
+                                sampler=SamplerConfig(temperature=temperature, top_p=top_p, top_k=top_k, presence_penalty=penalty),
                                 seed=4294967297, stop_token_ids=stop_ids, token_callback=lambda ids: emitted.extend(ids), capture_final_state=capture)
                             # Decode calls only: real prefill consumed [1] then [2].
                             calls = [{"ids": np.array(ids).reshape(-1).tolist(), "hidden": hidden} for ids, hidden in model.calls[2:]]
                             final = None if out.final_state is None else dict(shape=out.final_state.final_logits.shape,
                                 sha256=hashlib.sha256(np.array(out.final_state.final_logits).tobytes()).hexdigest())
-                            print(json.dumps(dict(pipeline=pipeline, asynchronous=asynchronous, temperature=temperature,
+                            print(json.dumps(dict(legacy_pipeline=pipeline, asynchronous=asynchronous, temperature=temperature, top_p=top_p, top_k=top_k,
                                 penalty=penalty, limit=limit, capture=capture, stops=sorted(stop_ids), tokens=out.tokens,
-                                emitted=emitted, calls=calls, modes=model.modes, final=final,
+                                emitted=emitted, calls=calls, final=final,
                                 verify_calls=out.stats.verify_calls, finish=out.finish_reason)), flush=True)
 
 

tools/mtplx-policy-reference.py

@@ -6,6 +6,8 @@ import os
 
 parser = argparse.ArgumentParser(description=__doc__)
 parser.add_argument("model")
+parser.add_argument("--server-only", action="store_true",
+    help="Compare server defaults without the older boundary-cache option API")
 args = parser.parse_args()
 print("reference: importing policy functions without model loading", flush=True)
 from mtplx.server.openai import _server_runtime_env_overrides
@@ -47,6 +49,10 @@ for mtp, env in itertools.product((False, True), cases):
     overrides = _server_runtime_env_overrides(argparse.Namespace(
         model=args.model, generation_mode="mtp" if mtp else "ar"), None)
     os.environ.update(overrides)
+    if args.server_only:
+        print(json.dumps(dict(event="policy_reference", mtp=mtp, environment=env,
+            overrides=overrides, server_only=True)), flush=True)
+        continue
     # Native app selects bounded prefill; emulate that request-local policy.
     os.environ["MTPLX_SUSTAINED_PREFILL"] = "1"
     error = None
@@ -59,7 +65,7 @@ for mtp, env in itertools.product((False, True), cases):
                 lazy_kv=_lazy_snapshot_enabled(), store=g._store_on_prefill_env_enabled(),
                 store_min=g._store_on_prefill_min_suffix(), fused_max=g._small_suffix_fused_max(),
                 boundaries=g._gdn_boundary_capture_enabled(), boundary_max=g._gdn_boundary_max_count(),
-                boundary_tail=g._gdn_boundary_tail_interval(), pipeline=g._env_truthy("MTPLX_AR_PIPELINE"),
+                boundary_tail=g._gdn_boundary_tail_interval(),
                 sync_eval=not g._env_truthy("MTPLX_ASYNC_AR") or bool(os.environ.get("MTPLX_EVAL_AUDIT")),
                 capture=g._family_capture_commit_enabled() if mtp else None,
                 lazy_history=g._env_truthy("MTPLX_LAZY_MTP_HISTORY_APPEND"),