Reject invalid Qwen logits and discard poisoned sessions

tools/issue105-reference.py

@@ -0,0 +1,36 @@
+"""Model-free #105 receipt from the pinned MTPLX/MLX oracle, never production."""
+import json
+import subprocess
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+from mtplx import generation
+from mtplx.sampling import NonFiniteLogitsError, SamplerConfig
+
+reference = Path(generation.__file__).resolve().parent.parent
+assert subprocess.check_output(["git", "-C", str(reference), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112"
+assert mx.__version__ == "0.32.2"
+rows = {
+    "nan": np.full(64, np.nan, np.float32),
+    "positive_inf": np.array([np.inf] + [0.] * 63, np.float32),
+    "all_negative_inf": np.full(64, -np.inf, np.float32),
+    "valid_mask": np.array([-np.inf, 2., -np.inf, 1.] + [-np.inf] * 60, np.float32),
+}
+for name, values in rows.items():
+    row = mx.array(values)
+    for temperature in (0., .6, 1.):
+        config = SamplerConfig(temperature=temperature, top_p=.95, top_k=20)
+        try:
+            token, _ = generation._sample_from_logits(row, config, np.random.default_rng(0))
+            result = {"token": token}
+        except NonFiniteLogitsError as error:
+            result = {"error": str(error), "nan": error.nan_count, "inf": error.inf_count, "vocab": error.vocab_size}
+        print(json.dumps(dict(case="ordinary", row=name, temperature=temperature, **result)), flush=True)
+    config = SamplerConfig(temperature=1., top_p=.95, top_k=20)
+    token = int(generation._mx_lazy_sample(row, config, mx.random.key(7)).item())
+    print(json.dumps(dict(case="lazy", row=name, token=token)), flush=True)
+    # The specialized MTP greedy helper still lacks the ordinary guard at this
+    # revision. Keep the conflict with #105 visible; never label this parity.
+    token, _ = generation._greedy_draft_token_and_top_values(row.reshape(1,1,-1), topk=2)
+    print(json.dumps(dict(case="mtp_greedy_specialized", row=name, token=token)), flush=True)