| """WIN2-R3 distributional gate: registered prompts, fixed-position probes, KL/top-1 compare, A/A band, gate (CPU only).""" |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import math |
| from collections.abc import Sequence |
| from pathlib import Path |
|
|
| import pytest |
|
|
| from offload.gb10_budget import window_cells as wc |
| from offload.launch.windows import dist_gate as dg |
| from offload.prefill_maxlen.evidence import Json, number, record |
|
|
| def materials() -> str: |
| from offload.prefill_trace.window_requests import HEADER |
| words = ' '.join(f'w{i % 997}' for i in range(2_600_000)) |
| return HEADER + words |
|
|
|
|
| class FakeDist: |
| """1 token = 4 characters; top-20 depends on the prefix only; `shift` moves that much probability from the top token |
| to the second (a candidate change); `flip` swaps the top-1 at positions whose prefix hash is divisible by `flip`.""" |
| def __init__(self, shift: float = 0.0, flip: int = 0) -> None: |
| self.shift, self.flip = shift, flip |
| self.probes = 0 |
| self.salts: set[str] = set() |
|
|
| def model(self) -> dict[str, Json]: |
| return {'id': 'fixture', 'max_model_len': wc.CAP} |
|
|
| def token_ids(self, text: str) -> list[int]: |
| return [ord(c) for c in text[::4]] |
|
|
| def top20(self, prefix: Sequence[int], salt: str) -> dict[str, Json]: |
| self.probes += 1 |
| self.salts.add(salt) |
| h = int(hashlib.sha256(json.dumps(list(prefix[-50:])).encode()).hexdigest()[:8], 16) |
| probs = [0.5 - self.shift, 0.3 + self.shift] + [0.19 / 18] * 18 |
| if self.flip and h % self.flip == 0: |
| probs[0], probs[1] = probs[1], probs[0] |
| return {f'token_id:{100 + (h + i) % 5000}': math.log(p) for i, p in enumerate(probs)} |
|
|
|
|
| def test_registered_prompt_set() -> None: |
| targets = [p.target for p in dg.DIST_PROMPTS] |
| assert len(dg.DIST_PROMPTS) == 64 and targets.count(1024) == 32 and targets.count(32000) == 24 |
| assert targets.count(150000) == 8 and len({p.key for p in dg.DIST_PROMPTS}) == 64 |
| assert len(dg.selected(['short', '32k'])) == 56 |
| with pytest.raises(ValueError): |
| dg.selected(['huge']) |
|
|
|
|
| @pytest.fixture(scope='module') |
| def run() -> dict[str, Json]: |
| text = materials() |
| prompts = dg.selected(['short'])[:4] + dg.selected(['32k'])[:2] |
| api = FakeDist() |
| cut = dg.cuts(api, text, {'x': 'y'}, prompts) |
| rows = record(cut['prompts']) |
| control = dg.probe(api, text, rows, 'W-c') |
| return {'materials': text, 'rows': rows, 'control': control, 'probes': api.probes} |
|
|
|
|
| def test_probe_records_eight_fixed_positions_without_text(run: dict[str, Json]) -> None: |
| raw = record(run['control'])['positions'] |
| assert isinstance(raw, list) |
| positions = [record(p) for p in raw] |
| assert len(positions) == 6 * dg.POSITIONS and run['probes'] == 6 * dg.POSITIONS |
| assert [p['k'] for p in positions[:8]] == list(range(8)) |
| assert 'w1' not in json.dumps(run['control']) |
|
|
|
|
| def test_compare_identical_and_shifted(run: dict[str, Json]) -> None: |
| text, rows = str(run['materials']), record(run['rows']) |
| control = record(run['control']) |
| same = dg.compare(control, dg.probe(FakeDist(), text, rows, 'W-x')) |
| assert same['mean_kl'] == 0.0 and same['top1_agree'] == same['positions'] == 48 |
| shifted = dg.compare(control, dg.probe(FakeDist(shift=0.08), text, rows, 'W-y')) |
| assert number(shifted['mean_kl']) > 0.01 and shifted['top1_agree'] == 48 |
| flipped = dg.compare(control, dg.probe(FakeDist(flip=3), text, rows, 'W-z')) |
| assert number(flipped['top1_agree']) < 48 |
| with pytest.raises(ValueError, match='different prompts'): |
| dg.compare(control, {'positions': record(dg.probe(FakeDist(), text, {'k': rows[next(iter(rows))]}, 's'))['positions']}) |
|
|
|
|
| def test_band_and_gate(run: dict[str, Json], tmp_path: Path) -> None: |
| text, rows = str(run['materials']), record(run['rows']) |
| control = record(run['control']) |
| aa = dg.compare(control, dg.probe(FakeDist(shift=0.002), text, rows, 'W-aa')) |
| band = dg.register_band([aa]) |
| assert number(band['kl_ceiling']) == pytest.approx(1.5 * number(band['aa_mean_kl_p99'])) and band['top1_min'] == 48 |
| assert dg.gate(dg.compare(control, dg.probe(FakeDist(shift=0.002), text, rows, 'W-n')), band)['verdict'] == 'PASS' |
| bad = dg.gate(dg.compare(control, dg.probe(FakeDist(shift=0.08), text, rows, 'W-b')), band) |
| assert bad['verdict'] == 'LOSS' and 'mean KL' in str(bad['reasons']) |
| flips = dg.gate(dg.compare(control, dg.probe(FakeDist(shift=0.002, flip=3), text, rows, 'W-f')), band) |
| assert flips['verdict'] == 'LOSS' and 'top-1' in str(flips['reasons']) |
| tampered = {**band, 'kl_ceiling': 10.0} |
| with pytest.raises(ValueError, match='changed after registration'): |
| dg.gate(aa, tampered) |
| with pytest.raises(ValueError, match='different position sets'): |
| dg.gate({**aa, 'positions': 8}, band) |
|
|
|
|
| def test_cli_compare_band_gate(run: dict[str, Json], tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: |
| text, rows = str(run['materials']), record(run['rows']) |
| paths = {} |
| for name, api in (('c', FakeDist()), ('aa', FakeDist(shift=0.002)), ('x', FakeDist(shift=0.08))): |
| paths[name] = tmp_path / f'{name}.json' |
| paths[name].write_text(json.dumps(dg.probe(api, text, rows, f'W-{name}'))) |
| assert dg.main(['compare', str(paths['c']), str(paths['aa']), '--output', str(tmp_path / 'aa-cmp.json')]) == 0 |
| assert dg.main(['band', str(tmp_path / 'aa-cmp.json'), '--output', str(tmp_path / 'band.json')]) == 0 |
| assert dg.main(['compare', str(paths['c']), str(paths['x']), '--output', str(tmp_path / 'x-cmp.json')]) == 0 |
| assert dg.main(['gate', str(tmp_path / 'x-cmp.json'), '--band', str(tmp_path / 'band.json')]) == 2 |
| assert dg.main(['probe', '--cuts', 'x', '--salt', 's', '--output', str(tmp_path / 'p.json')]) == 0 |
| monkeypatch.delenv('WINDOW_COORDINATOR', raising=False) |
| assert dg.main(['probe', '--execute', '--cuts', 'x', '--salt', 's', '--output', str(tmp_path / 'p.json')]) == 2 |
|
|
|
|
| def test_top1_test_is_exact_and_calibrated() -> None: |
| assert dg.top1_p(0, 5, 1) == 1.0 |
| assert dg.top1_p(7, 0, 1) == pytest.approx(0.5 ** 7) |
| assert dg.top1_p(6, 0, 1) == pytest.approx(0.5 ** 6) |
| assert dg.top1_p(5, 3, 1) > 0.1 |
| assert dg.top1_p(4, 0, 3) == pytest.approx(0.25 ** 4) |
| rng = __import__('random').Random(5) |
| rate, trials, losses = 3 / 512, 3000, 0 |
| for _ in range(trials): |
| aa = sum(rng.random() < rate for _ in range(512)) |
| cand = sum(rng.random() < rate for _ in range(512)) |
| losses += dg.top1_p(cand, aa, 1) < dg.TOP1_ALPHA |
| assert losses / trials < 0.012 |
| with pytest.raises(ValueError): |
| dg.top1_p(-1, 0, 1) |
|
|
|
|
| def test_low_common_mass_positions_can_lose(run: dict[str, Json]) -> None: |
| text, rows = str(run['materials']), record(run['rows']) |
| control = record(run['control']) |
| aa = dg.compare(control, dg.probe(FakeDist(shift=0.002), text, rows, 'W-aa2')) |
| band = dg.register_band([aa]) |
| assert band['aa_low_common_mass'] == 0 |
| shifted = dict(dg.compare(control, dg.probe(FakeDist(shift=0.002), text, rows, 'W-lc'))) |
| shifted['low_common_mass_positions'] = 9 |
| g = dg.gate(shifted, band) |
| assert g['verdict'] == 'LOSS' and 'common mass' in str(g['reasons']) |
|
|
