hugo/DS4Server
1
0 Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

Reject invalid Qwen logits and discard poisoned sessions

Browse Source
This commit is contained in:
Georg Bauer committed 2026-10-04 13:29:09 +02:00
1 parent 1de78fa3ef
commit 5ecf8ab4f3
31 files changed
+1812 -746
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tools/mtplx-kernel-fixtures.py
+25 -18
View File
@@ -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)