| """WIN2-R3 distributional gate for candidates (F-METHOD §3.4): is the candidate's next-token distribution the same model? |
| |
| Probe: 64 public-corpus prompts (32 short ~1K tokens, 24 at 32K, 8 at 150K; offsets registered below), each tokenized |
| once (/tokenize, raw text), then the top-20 next-token logprobs at 8 FIXED positions: the prompt's last 8 tokens are the |
| fixed continuation, and position k asks for the next token after ids[:L - 8 + k] (/v1/completions, max_tokens 1, |
| temperature 0, logprobs 20, token-id prompts, one cache salt per run so only the first probe of a prompt prefills). No |
| generation, so no generated text exists; only token ids and logprob vectors are recorded, on Controller (never git). |
| |
| Compare (control || candidate) per position with greedy_kl.top20_kl (shared top-20 + OTHER; common mass >= 0.99 is |
| reported per position) and top-1 agreement. Gate against an A/A band registered BEFORE use from two production starts: |
| - kl_ceiling = 1.5 x the 99th percentile of the A/A MEAN KL (seeded bootstrap over the A/A positions); LOSS when the |
| candidate's mean KL exceeds it (simulated no-effect false LOSS 0.0-0.2 %); |
| - top-1: F-METHOD's literal rule (LOSS when agreement < A/A minimum - 1) rejects a no-effect candidate 13-34 % of the |
| time once the A/A itself shows >= 1 expected near-tie flip (one A/A run is a noisy minimum). The gate therefore uses an |
| exact conditional test of the two flip counts instead: with D_c candidate and D_a A/A disagreements over k registered |
| A/A comparisons, D_c | D_c + D_a ~ Binomial(D_c + D_a, 1 / (k + 1)) under equal flip rates; LOSS when the one-sided |
| p < 0.01 (false LOSS <= 1 % by construction). The literal rule is still reported (`fmethod_top1_rule`). |
| - common mass: positions whose shared top-20 mass is < 0.99 (where the coarse KL is not trustworthy) are counted; LOSS |
| when the candidate has significantly more of them than the A/A (same exact test). Greedy identity stays the restore |
| gate (W8), unchanged. COST (ESTIMATED from MEASURED TTFTs): one cold prefill per prompt per arm, 32 x ~0.5 s + 24 x ~15 s |
| + 8 x ~69 s ~ 15.5 min per arm; `--subset short,32k` drops the 150K prompts (~6.5 min). |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import math |
| import os |
| import random |
| import statistics |
| import sys |
| import time |
| from collections.abc import Mapping, Sequence |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Protocol, cast |
|
|
| from offload.gb10_budget import window_cells as wc |
| from offload.launch.windows import greedy_kl |
| from offload.prefill_maxlen.evidence import Json, number, record |
|
|
| SCHEMA = 'dsv41-dist-gate-v1' |
| POSITIONS = 8 |
| KL_FACTOR = 1.5 |
| BOOTSTRAP = 2000 |
| BOOTSTRAP_SEED = 20261027 |
| TOP1_ALPHA = 0.01 |
| COORDINATOR = 'Deepseek/coordinator' |
|
|
|
|
| @dataclass(frozen=True) |
| class DistPrompt: |
| key: str |
| target: int |
| offset: int |
|
|
|
|
| def _prompts() -> tuple[DistPrompt, ...]: |
| out = [DistPrompt(f'dist-short-{i:02d}', 1024, 200_000 + 10_000 * i) for i in range(32)] |
| out += [DistPrompt(f'dist-32k-{i:02d}', 32000, 1_000_000 + 120_000 * i) for i in range(24)] |
| out += [DistPrompt(f'dist-150k-{i}', 150000, 4_200_000 + 560_000 * i) for i in range(8)] |
| return tuple(out) |
|
|
|
|
| DIST_PROMPTS = _prompts() |
| SUBSETS = {'short': 1024, '32k': 32000, '150k': 150000} |
|
|
|
|
| def selected(subset: Sequence[str] | None) -> tuple[DistPrompt, ...]: |
| if not subset: |
| return DIST_PROMPTS |
| unknown = set(subset) - set(SUBSETS) |
| if unknown: |
| raise ValueError(f'unknown prompt subsets {sorted(unknown)}; choose from {sorted(SUBSETS)}') |
| targets = {SUBSETS[s] for s in subset} |
| return tuple(p for p in DIST_PROMPTS if p.target in targets) |
|
|
|
|
| class DistApi(Protocol): |
| def model(self) -> dict[str, Json]: ... |
| def token_ids(self, text: str) -> list[int]: ... |
| def top20(self, prefix: Sequence[int], salt: str) -> dict[str, Json]: ... |
|
|
|
|
| class Http(wc.Http): |
| def token_ids(self, text: str) -> list[int]: |
| r = self.post('/tokenize', {'model': self.model()['id'], 'prompt': text, 'add_special_tokens': True}) |
| ids = r.get('tokens') |
| if not isinstance(ids, list) or not all(isinstance(t, int) and not isinstance(t, bool) for t in ids): |
| raise ValueError('tokenize returned no token ids') |
| return cast(list[int], ids) |
|
|
| def top20(self, prefix: Sequence[int], salt: str) -> dict[str, Json]: |
| r = self.post('/v1/completions', {'model': self.model()['id'], 'prompt': cast(Json, list(prefix)), 'max_tokens': 1, |
| 'temperature': 0, 'logprobs': 20, 'return_tokens_as_token_ids': True, |
| 'priority': 10, 'cache_salt': salt}) |
| choices = r.get('choices') |
| if not isinstance(choices, list) or len(choices) != 1: |
| raise ValueError('one completion choice expected') |
| tops = record(record(choices[0])['logprobs'])['top_logprobs'] |
| if not isinstance(tops, list) or len(tops) != 1: |
| raise ValueError('one top-logprob position expected') |
| return record(tops[0]) |
|
|
|
|
| def cuts(api: DistApi, materials: str, hashes: Mapping[str, Json], prompts: Sequence[DistPrompt]) -> dict[str, Json]: |
| """Numeric cuts (characters, tokens, sha256) of the registered prompts; built once, reused by every arm.""" |
| rows: dict[str, Json] = {} |
| for p in prompts: |
| chars, tokens = wc.cut(materials, wc.Prompt(p.key, p.target, p.offset), lambda t: len(api.token_ids(t))) |
| text = wc.render(materials, p.offset, chars) |
| rows[p.key] = {'target': p.target, 'offset': p.offset, 'chars': chars, 'tokens': tokens, |
| 'sha256': hashlib.sha256(text.encode()).hexdigest()} |
| return {'schema': SCHEMA, 'type': 'cuts', 'corpus_sha256': cast(Json, dict(hashes)), 'prompts': rows} |
|
|
|
|
| def probe(api: DistApi, materials: str, cut_rows: Mapping[str, Json], salt_prefix: str) -> dict[str, Json]: |
| """Top-20 at the 8 fixed positions of every prompt in `cut_rows` on the server under test.""" |
| if not salt_prefix.strip(): |
| raise ValueError('a cache salt prefix per run is required') |
| positions: list[Json] = [] |
| for key, value in record(cast(Json, dict(cut_rows))).items(): |
| row = record(value) |
| text = wc.render(materials, int(number(row['offset'])), int(number(row['chars']))) |
| if hashlib.sha256(text.encode()).hexdigest() != row.get('sha256'): |
| raise ValueError(f'{key}: rebuilt prompt differs from its registered cut') |
| ids = api.token_ids(text) |
| if len(ids) <= POSITIONS: |
| raise ValueError(f'{key}: prompt shorter than the fixed continuation') |
| salt = f'{salt_prefix}:dist:{key}' |
| for k in range(POSITIONS): |
| at = len(ids) - POSITIONS + k |
| positions.append({'prompt': key, 'position': at, 'k': k, |
| 'prefix_sha256': hashlib.sha256(json.dumps(ids[:at]).encode()).hexdigest(), |
| 'top20': api.top20(ids[:at], salt)}) |
| return {'schema': SCHEMA, 'type': 'probe', 'positions': positions, 'utc': time.strftime('%Y-%m-%dT%H:%M:%SZ')} |
|
|
|
|
| def _top1(top: Mapping[str, Json]) -> str: |
| return max(top, key=lambda k: number(top[k])) |
|
|
|
|
| def compare(control: Mapping[str, Json], candidate: Mapping[str, Json]) -> dict[str, Json]: |
| a, b = [record(p) for p in cast(list[Json], control['positions'])], [record(p) for p in cast(list[Json], candidate['positions'])] |
| if [(p['prompt'], p['position'], p['prefix_sha256']) for p in a] != [(p['prompt'], p['position'], p['prefix_sha256']) for p in b]: |
| raise ValueError('control and candidate probed different prompts or positions') |
| per: list[Json] = [] |
| kls: list[float] = [] |
| agree = low = 0 |
| for x, y in zip(a, b): |
| r = greedy_kl.top20_kl(record(x['top20']), record(y['top20'])) |
| same = _top1(record(x['top20'])) == _top1(record(y['top20'])) |
| agree += same |
| low += r['common_mass'] < greedy_kl.MIN_COMMON_MASS |
| kls.append(r['kl_nats']) |
| per.append({'prompt': x['prompt'], 'k': x['k'], 'kl_nats': r['kl_nats'], 'common_mass': r['common_mass'], |
| 'top1_agree': same}) |
| return {'schema': SCHEMA, 'type': 'compare', 'positions': len(kls), 'mean_kl': statistics.fmean(kls), |
| 'max_kl': max(kls), 'top1_agree': agree, 'low_common_mass_positions': low, 'per_position': per} |
|
|
|
|
| def _p99(values: Sequence[float]) -> float: |
| ordered = sorted(values) |
| return ordered[min(len(ordered) - 1, math.ceil(0.99 * len(ordered)) - 1)] |
|
|
|
|
| def register_band(aa: Sequence[Mapping[str, Json]], *, factor: float = KL_FACTOR) -> dict[str, Json]: |
| """A/A band from compare() results of production vs production (two starts). Register it before any candidate use.""" |
| if not aa: |
| raise ValueError('at least one A/A comparison is required') |
| counts = {int(number(c['positions'])) for c in aa} |
| if len(counts) != 1: |
| raise ValueError('A/A comparisons cover different position sets') |
| rng = random.Random(BOOTSTRAP_SEED) |
| means: list[float] = [] |
| for c in aa: |
| kls = [number(record(p)['kl_nats']) for p in cast(list[Json], c['per_position'])] |
| means += [statistics.fmean(rng.choices(kls, k=len(kls))) for _ in range(BOOTSTRAP)] |
| band: dict[str, Json] = {'schema': SCHEMA, 'type': 'band', 'positions': counts.pop(), 'factor': factor, |
| 'aa_mean_kl_p99': _p99(means), 'kl_ceiling': factor * _p99(means), |
| 'top1_min': min(int(number(c['top1_agree'])) for c in aa), 'aa_comparisons': len(aa), |
| 'aa_top1_disagreements': sum(int(number(c['positions']) - number(c['top1_agree'])) for c in aa), |
| 'aa_low_common_mass': sum(int(number(c['low_common_mass_positions'])) for c in aa), |
| 'bootstrap': {'resamples': BOOTSTRAP, 'seed': BOOTSTRAP_SEED}, |
| 'registered_utc': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())} |
| band['sha256'] = hashlib.sha256(json.dumps(band, sort_keys=True).encode()).hexdigest() |
| return band |
|
|
|
|
| def gate(result: Mapping[str, Json], band: Mapping[str, Json]) -> dict[str, Json]: |
| body = {k: v for k, v in band.items() if k != 'sha256'} |
| if band.get('sha256') != hashlib.sha256(json.dumps(body, sort_keys=True).encode()).hexdigest(): |
| raise ValueError('A/A band changed after registration') |
| if result.get('positions') != band.get('positions'): |
| raise ValueError('candidate comparison and A/A band cover different position sets') |
| reasons = [] |
| if number(result['mean_kl']) > number(band['kl_ceiling']): |
| reasons.append(f'mean KL {number(result["mean_kl"]):.3g} > ceiling {number(band["kl_ceiling"]):.3g}') |
| flips = int(number(result['positions']) - number(result['top1_agree'])) |
| p = top1_p(flips, int(number(band['aa_top1_disagreements'])), int(number(band['aa_comparisons']))) |
| if p < TOP1_ALPHA: |
| reasons.append(f'top-1 disagreements {flips} vs A/A {band["aa_top1_disagreements"]} over ' |
| f'{band["aa_comparisons"]} comparison(s): one-sided p {p:.4f} < {TOP1_ALPHA}') |
| low = int(number(result['low_common_mass_positions'])) |
| p_low = top1_p(low, int(number(band['aa_low_common_mass'])), int(number(band['aa_comparisons']))) |
| if p_low < TOP1_ALPHA: |
| reasons.append(f'positions with top-20 common mass < {greedy_kl.MIN_COMMON_MASS} {low} vs A/A ' |
| f'{band["aa_low_common_mass"]}: one-sided p {p_low:.4f} < {TOP1_ALPHA}') |
| return {'verdict': 'LOSS' if reasons else 'PASS', 'reasons': cast(Json, reasons), 'band_sha256': band['sha256'], |
| 'mean_kl': result['mean_kl'], 'top1_agree': result['top1_agree'], 'top1_disagreements': flips, |
| 'top1_p': p, 'fmethod_top1_rule': number(result['top1_agree']) < number(band['top1_min']) - 1, |
| 'low_common_mass_positions': result['low_common_mass_positions']} |
|
|
|
|
| def top1_p(candidate_flips: int, aa_flips: int, aa_comparisons: int) -> float: |
| """P(X >= candidate_flips), X ~ Binomial(candidate_flips + aa_flips, 1 / (aa_comparisons + 1)): the exact |
| conditional test that the candidate's top-1 flip rate exceeds the A/A rate (Poisson counts, same positions).""" |
| if candidate_flips < 0 or aa_flips < 0 or aa_comparisons < 1: |
| raise ValueError('invalid flip counts') |
| n, q = candidate_flips + aa_flips, 1 / (aa_comparisons + 1) |
| return sum(math.comb(n, i) * q ** i * (1 - q) ** (n - i) for i in range(candidate_flips, n + 1)) |
|
|
|
|
| def _read(path: Path) -> dict[str, Json]: |
| return record(cast(Json, json.loads(path.read_text()))) |
|
|
|
|
| def _write(path: Path, data: Mapping[str, Json]) -> None: |
| with path.open('x') as handle: |
| handle.write(json.dumps(data, allow_nan=False) + '\n') |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__) |
| sub = ap.add_subparsers(dest='cmd', required=True) |
| for name in ('cuts', 'probe'): |
| s = sub.add_parser(name) |
| s.add_argument('--execute', action='store_true') |
| s.add_argument('--url', default=wc.URL) |
| s.add_argument('--data-dir', type=Path, default=Path('/opt/dsv41/data/bench-data')) |
| s.add_argument('--output', type=Path, required=True) |
| s.add_argument('--subset', help='comma list of short,32k,150k (default: all 64 prompts)') |
| if name == 'probe': |
| s.add_argument('--cuts', type=Path, required=True) |
| s.add_argument('--salt', required=True) |
| c = sub.add_parser('compare') |
| c.add_argument('control', type=Path) |
| c.add_argument('candidate', type=Path) |
| c.add_argument('--output', type=Path, required=True) |
| b = sub.add_parser('band') |
| b.add_argument('aa', type=Path, nargs='+', help='compare() outputs of production vs production') |
| b.add_argument('--output', type=Path, required=True) |
| g = sub.add_parser('gate') |
| g.add_argument('comparison', type=Path) |
| g.add_argument('--band', type=Path, required=True) |
| args = ap.parse_args(argv) |
| if args.cmd in ('cuts', 'probe'): |
| prompts = selected(args.subset.split(',') if args.subset else None) |
| if not args.execute: |
| print(json.dumps({'executed': False, 'prompts': len(prompts), 'positions': POSITIONS * len(prompts)})) |
| return 0 |
| if os.environ.get('WINDOW_COORDINATOR') != COORDINATOR: |
| print(f'STOP: probes are coordinator-only (WINDOW_COORDINATOR={COORDINATOR})', file=sys.stderr) |
| return 2 |
| from offload.prefill_trace.window_requests import corpus |
| materials, hashes = corpus(args.data_dir) |
| api = Http(args.url) |
| if args.cmd == 'cuts': |
| _write(args.output, cuts(api, materials, hashes, prompts)) |
| else: |
| registered = _read(args.cuts) |
| if registered.get('corpus_sha256') != hashes: |
| raise ValueError('public corpus changed since the cuts were built') |
| rows = record(registered['prompts']) |
| _write(args.output, probe(api, materials, {p.key: rows[p.key] for p in prompts}, args.salt)) |
| return 0 |
| if args.cmd == 'compare': |
| _write(args.output, compare(_read(args.control), _read(args.candidate))) |
| return 0 |
| if args.cmd == 'band': |
| _write(args.output, register_band([_read(p) for p in args.aa])) |
| return 0 |
| verdict = gate(_read(args.comparison), _read(args.band)) |
| print(json.dumps(verdict, allow_nan=False)) |
| return 0 if verdict['verdict'] == 'PASS' else 2 |
|
|
|
|
| if __name__ == '__main__': |
| raise SystemExit(main()) |
|
|
