| |
| """Compare original baseline and final kernel on identical inputs and fp64 truth.""" |
| import importlib.util |
| import json |
| from pathlib import Path |
| import sys |
|
|
| import torch |
|
|
|
|
| def load(path, name): |
| spec = importlib.util.spec_from_file_location(name, path) |
| module = importlib.util.module_from_spec(spec) |
| sys.modules[name] = module |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| kind = sys.argv[1] |
| release_root = Path(__file__).resolve().parent |
| here = release_root / kind |
| bench = load(here / 'benchmark.py', f'{kind}_benchmark') |
| baseline = load(here / 'baseline.py', f'{kind}_baseline') |
| candidate = load(release_root / 'build/torch-cuda' / f'_{kind}.py', f'{kind}_candidate') |
| scope = {} |
| exec(json.loads((here / 'definition.json').read_text())['reference'], scope) |
| reference = scope['run'] |
| device = torch.device('cuda') |
| report = {'component': kind, 'gpu': torch.cuda.get_device_name(), 'torch': torch.__version__, |
| 'method': 'Identical inputs for both implementations, original judge unchanged, fp64 rebuilt truth.', |
| 'global_seed_plus_batch': 20261026, 'relative_tolerance': bench.REL_TOL, |
| 'absolute_cap_factor': bench.CAP, 'reference_error_floor_factor': bench.FLOOR, |
| 'individual_state_check_is_additional': True, 'rows': []} |
| for batch in (16, 32, 64, 128, 256, 512): |
| torch.manual_seed(20261026 + batch) |
| first = bench.structured_inputs(batch, device, seed=batch) |
| ref1 = reference(*[x.clone() for x in first]) |
| second = (*ref1, *bench.update_buffer(batch, device, torch.Generator(device=device).manual_seed(batch + 1)), first[-1]) |
| random = bench.random_inputs(batch, device, seed=batch + 2) |
| for family, inputs, ref in (('structured', first, ref1), ('second_window', second, None), ('random', random, None)): |
| if ref is None: |
| ref = reference(*[x.clone() for x in inputs]) |
| exact = bench.exact_state(inputs) |
| ref_error = (bench.decode(*ref) - exact).abs().mean((-2, -1)) |
| magnitude = exact.abs().mean((-2, -1)) |
| tight = (ref_error >= bench.FLOOR * magnitude) & (ref_error > 0) |
| row = {'batch_size': batch, 'family': family} |
| for name, impl in (('baseline', baseline), ('current', candidate)): |
| output = impl.run(*[x.clone() for x in inputs]) |
| torch.cuda.synchronize() |
| passed, ratio = bench.judge(output, ref, exact) |
| error = (bench.decode(*output) - exact).abs().mean((-2, -1)) |
| strict = torch.where(tight, error <= (1 + bench.REL_TOL) * ref_error, error <= bench.CAP * magnitude) |
| row[name] = { |
| 'original_judge_passed': passed, |
| 'worst_head_error_ratio': float(ratio), |
| 'max_individual_head_error_ratio': float((error[tight] / ref_error[tight]).max()), |
| 'individual_states_checked': int(strict.numel()), |
| 'individual_states_passing_same_bound': int(strict.sum()), |
| 'max_state_relative_mae': float((error[magnitude > 0] / magnitude[magnitude > 0]).max()), |
| } |
| assert passed, (kind, batch, family, name, ratio) |
| report['rows'].append(row) |
| (here / 'precision_comparison.json').write_text(json.dumps(report, indent=2) + '\n') |
| print(kind, batch, 'baseline / current:', [(r['family'], r['baseline']['worst_head_error_ratio'], r['current']['worst_head_error_ratio']) for r in report['rows'][-3:]], flush=True) |
| report['status'] = 'original_benchmark_pass' |
| report['summary'] = {} |
| for name in ('baseline', 'current'): |
| values = [row[name] for row in report['rows']] |
| report['summary'][name] = { |
| 'original_cases_passed': sum(v['original_judge_passed'] for v in values), |
| 'worst_original_benchmark_error_ratio': max(v['worst_head_error_ratio'] for v in values), |
| 'individual_states_checked': sum(v['individual_states_checked'] for v in values), |
| 'additional_individual_state_failures': sum(v['individual_states_checked'] - v['individual_states_passing_same_bound'] for v in values), |
| 'max_individual_head_error_ratio': max(v['max_individual_head_error_ratio'] for v in values), |
| } |
| (here / 'precision_comparison.json').write_text(json.dumps(report, indent=2) + '\n') |
| print(kind, 'complete', flush=True) |
|
|
