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)

tools/mtplx-kernel-fixtures.py

@@ -1396,7 +1396,7 @@ def ple_connected_fixtures():
 def batched_sampler_fixtures():
     """Both actual batch routes, their whole-batch fallback and public methods."""
     from mtplx import fast_sampling as fs
-    from mtplx.sampling import SamplerConfig
+    from mtplx.sampling import SamplerConfig, NonFiniteLogitsError
     assert mx.__version__ == "0.32.2" and np.__version__ == "2.4.4"
     def digest(values, dtype):
         return hashlib.sha256(np.asarray(values, dtype=dtype).tobytes()).hexdigest()
@@ -1407,13 +1407,13 @@ def batched_sampler_fixtures():
         samples = []
         for i in range(32):
             try: samples.append(batch.sample(i % len(batch.token_ids), rng))
-            except ValueError: samples.append(None)
+            except (ValueError,NonFiniteLogitsError): samples.append(None)
         converted = []
         for row in range(len(batch.token_ids)):
             try:
                 p = batch.to_distribution(row)
                 converted.append(dict(ids=digest(p.token_ids, '<i8'), probs=digest(p.probs, '<f8')))
-            except ValueError: converted.append(dict(error=True))
+            except (ValueError,NonFiniteLogitsError): converted.append(dict(error=True))
         lookups = []
         for row in (0, -1, len(batch.token_ids), -len(batch.token_ids)-1):
             for token in (-1, 0, 3, batch.vocab_size-1, batch.vocab_size):
@@ -1457,7 +1457,7 @@ def batched_sampler_fixtures():
                                         batch = (fs.bind_batched_top_k_distributions(config, vocab_size=vocab)(row)
                                             if route == "bound" else fs.batched_sparse_distributions_from_mlx_logits(row, config))
                                         result = receipt(batch)
-                                    except (ValueError, FloatingPointError): result = dict(error=True)
+                                    except (ValueError, FloatingPointError, NonFiniteLogitsError): result = dict(error=True)
                                     print(json.dumps(dict(vocab=vocab, dtype=str(dtype).split('.')[-1], layout=layout, mode=mode,
                                         temperature=temperature, top_p=top_p, k=k, route=route, result=result,
                                         evals=trace.copy(), host_calls=len(host_calls)), separators=(',',':')), flush=True)
@@ -1470,7 +1470,7 @@ def batched_sampler_fixtures():
                 ("nonfinite", [[0,1]], [[np.inf,1]]),
                 ("mismatch", [[0,1]], [[1,2,3]])):
                 try: result=receipt(fs.BatchedSparseDistributions(np.array(ids), np.array(probs), vocab_size=17))
-                except ValueError: result=dict(error=True)
+                except (ValueError,NonFiniteLogitsError): result=dict(error=True)
                 print(json.dumps(dict(constructor=name, result=result), separators=(',',':')), flush=True)
     finally:
         mx.eval, fs._host_sparse_distribution = original_eval, original_host
@@ -1781,7 +1781,7 @@ def resident_sampler_fixtures():
 def sampler_entry_fixtures():
     """Actual MTPLX generation entry, three consecutive draws and count updates."""
     from mtplx import generation
-    from mtplx.sampling import SamplerConfig, SparseDistribution
+    from mtplx.sampling import SamplerConfig, SparseDistribution, NonFiniteLogitsError
     assert mx.__version__ == "0.32.2"
     original_eval, original_argmax = mx.eval, mx.argmax
     trace = []
@@ -1825,7 +1825,7 @@ def sampler_entry_fixtures():
                                             result["probabilities"]=hashlib.sha256(probs.astype('<f8').tobytes()).hexdigest()
                                             if isinstance(p,SparseDistribution): result["ids"]=hashlib.sha256(p.token_ids.astype('<i8').tobytes()).hexdigest()
                                         counts[token]=counts.get(token,0)+1
-                                    except (ValueError,IndexError): result=dict(error=True)
+                                    except (ValueError,IndexError,NonFiniteLogitsError): result=dict(error=True)
                                     result["trace"]=list(trace); steps.append(result)
                                 print(json.dumps(dict(vocab=vocab,dtype=str(dtype).split('.')[-1],layout=layout,mode=mode,
                                     profile=profile,penalties=penalties,steps=steps,counts=list(counts.items()),
@@ -1871,7 +1871,7 @@ def sampler_penalty_fixtures():
 
 def sampler_mixed_distribution_fixtures():
     """Original dense/sparse choice, residual, and verify including RNG errors."""
-    from mtplx.sampling import SparseDistribution, sample_from_distribution, residual_distribution, verify_one_token
+    from mtplx.sampling import NonFiniteLogitsError, SparseDistribution, sample_from_distribution, residual_distribution, verify_one_token
     assert np.__version__ == "2.4.4"
     def bits(value): return int(np.asarray(value,dtype='<f8').view('<u8'))
     def digest(values,dtype): return hashlib.sha256(np.asarray(values,dtype=dtype).tobytes()).hexdigest()
@@ -1893,11 +1893,15 @@ def sampler_mixed_distribution_fixtures():
                         count = min(20,vocab)
                         ids = np.arange(count,dtype=np.int64)*(vocab-1)//max(1,count-1)
                         return SparseDistribution(ids,raw[ids],vocab) if kind == "s" else raw/raw.sum()
-                    target,draft = make(p,kinds[0]),make(q,kinds[1])
+                    try:
+                        target,draft = make(p,kinds[0]),make(q,kinds[1])
+                    except NonFiniteLogitsError:
+                        print(json.dumps(dict(vocab=vocab,mode=mode,kinds=kinds,construction_error=True)),flush=True)
+                        continue
                     rng = np.random.default_rng(12345)
                     try:
                         sampled = dict(tokens=[sample_from_distribution(target,rng) for _ in range(32)])
-                    except ValueError: sampled = dict(error=True)
+                    except (ValueError,NonFiniteLogitsError): sampled = dict(error=True)
                     sampled["next"] = bits(rng.random())
                     try:
                         r = residual_distribution(target,draft)
@@ -1905,7 +1909,7 @@ def sampler_mixed_distribution_fixtures():
                         residual = dict(kind="s" if sparse else "d",borrowed=r is target,
                             probabilities=digest(r.probs if sparse else r,'<f8'))
                         if sparse: residual["ids"] = digest(r.token_ids,'<i8')
-                    except ValueError: residual = dict(error=True)
+                    except (ValueError,NonFiniteLogitsError): residual = dict(error=True)
                     rng = np.random.default_rng(43210)
                     decisions=[]
                     for i in range(32):
@@ -1913,7 +1917,7 @@ def sampler_mixed_distribution_fixtures():
                         try:
                             decision=verify_one_token(target,draft,token,rng)
                             decisions.append([decision.accepted,decision.token_id,bits(decision.accept_probability)])
-                        except ValueError: decisions.append(None)
+                        except (ValueError,NonFiniteLogitsError): decisions.append(None)
                     print(json.dumps(dict(vocab=vocab,mode=mode,kinds=kinds,sampled=sampled,residual=residual,
                         decisions=decisions,next_verify=bits(rng.random())),separators=(',',':')),flush=True)
 
@@ -1959,24 +1963,27 @@ def sampler_argmax_fixtures():
 
 def sampler_fallback_fixtures():
     from mtplx.fast_sampling import sparse_distributions_from_mlx_logits
-    from mtplx.sampling import SamplerConfig, distribution_from_logits
+    from mtplx.sampling import SamplerConfig, distribution_from_logits, NonFiniteLogitsError
     assert mx.__version__ == "0.32.2"
     for vocab in (17, 4097, 248320):
         for mode in ("unique", "sharp", "equal", "nonfinite", "mixed"):
             values = sampler_serial_values(vocab, 4, mode)
             for top_p in (0.95, 1.0):
                 config = SamplerConfig(temperature=0.6, top_p=top_p, top_k=20)
-                distributions = sparse_distributions_from_mlx_logits(mx.array(values).reshape(4,vocab),config)
-                print(json.dumps(dict(kind="sparse",vocab=vocab,mode=mode,top_p=top_p,
-                    ids=[d.token_ids.tolist() for d in distributions],
-                    probabilities=[d.probs.astype('<f8').view('<u8').tolist() for d in distributions]),separators=(',',':')),flush=True)
+                try:
+                    distributions = sparse_distributions_from_mlx_logits(mx.array(values).reshape(4,vocab),config)
+                    result = dict(ids=[d.token_ids.tolist() for d in distributions],
+                        probabilities=[d.probs.astype('<f8').view('<u8').tolist() for d in distributions])
+                except NonFiniteLogitsError:
+                    result = dict(error=True)
+                print(json.dumps(dict(kind="sparse",vocab=vocab,mode=mode,top_p=top_p,**result),separators=(',',':')),flush=True)
                 for top_k in (0,20):
                     for temperature in (0.0,0.6):
                         try:
                             probs = distribution_from_logits(values[:vocab].astype(np.float64),
                                 SamplerConfig(temperature=temperature,top_p=top_p,top_k=top_k))
                             result = dict(sha256=hashlib.sha256(probs.astype('<f8').tobytes()).hexdigest())
-                        except ValueError:
+                        except (ValueError,NonFiniteLogitsError):
                             result = dict(error=True)
                         print(json.dumps(dict(kind="dense",vocab=vocab,mode=mode,top_p=top_p,
                             top_k=top_k,temperature=temperature,**result),separators=(',',':')),flush=True)