| |
| """Independent CPU arithmetic/packing versus actual v6 GPU recurrence. |
| |
| Strict dyadic fixtures require bitwise output, packed bytes and FP16 scales. |
| Random tests require exact kernel partition replay and every-token dense carry. |
| Dense-emulation readout is checked against an analytical FP32/FP16 error bound; |
| the independent CPU one-token comparison retains its fixed0.002 relative-L2 bound. |
| CPU-only preparation has a distinct receipt format and cannot pass a GPU gate. |
| """ |
| from __future__ import annotations |
| import argparse |
| import hashlib |
| import json |
| import inspect |
| import re |
| from pathlib import Path |
| import sys |
| import time |
|
|
| import torch |
| import triton |
| import triton.language as tl |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| PROTOCOL_SHA = '86d4e8dc85d79c2c867ce15a846c5938893c86d27e67ad174472975e654e9bff' |
| POLICIES = [('legacy', 'legacy', 1., None), ('stored_scale', 'stored_scale', 1., None), |
| ('clip4_095', 'stored_scale', .95, None), ('clip4_090', 'stored_scale', .90, None), |
| ('clip4_080', 'stored_scale', .80, None), ('prune_only', 'legacy', 1., 'prune_only'), |
| ('quant_only', 'legacy', 1., 'quant_only'), ('legacy_emulation', 'legacy', 1., 'legacy_emulation')] |
|
|
|
|
| def sha(path): return hashlib.sha256(Path(path).read_bytes()).hexdigest() |
|
|
|
|
| def half_away_cpu(value): |
| |
| absolute = value.abs() |
| floor = absolute.floor() |
| return value.sign()*(floor+((absolute-floor) >= .5).to(value.dtype)) |
|
|
|
|
| def pack_cpu(q): |
| result = torch.zeros((*q.shape[:-1], q.shape[-1]//2), dtype=torch.uint8) |
| for coordinate in range(q.shape[-1]): |
| nibble = (q[..., coordinate].to(torch.int64) & 15).to(torch.uint8) |
| result[..., coordinate//2] |= nibble << (4*(coordinate%2)) |
| return result |
|
|
|
|
| def decode_cpu(buffers): |
| shape = (*buffers['s8'].shape, 128) |
| state = torch.zeros(shape, dtype=torch.float32) |
| for coordinate in range(16): |
| shift = 4*(coordinate%2) |
| lo = (buffers['lo'][..., coordinate//2].int() >> shift) & 15 |
| hi = (buffers['hi'][..., coordinate//2].int() >> shift) & 15 |
| q = (hi << 4) | lo |
| q = torch.where(q > 127, q-256, q) |
| state[..., coordinate] = q.float()*buffers['s8'].float() |
| for coordinate in range(64): |
| q = (buffers['q4'][..., coordinate//2].int() >> (4*(coordinate%2))) & 15 |
| q = torch.where(q > 7, q-16, q) |
| state[..., 16+coordinate] = q.float()*buffers['s4'].float() |
| return state |
|
|
|
|
| def quantize_cpu(updated, mode, clip): |
| integers, scales = [], [] |
| for values, bound, factor in ((updated[..., :16], 127, 1.), (updated[..., 16:80], 7, clip)): |
| denominator = (values.abs().amax(-1)*factor/bound).clamp_min(1e-8) |
| stored = denominator.half() |
| actual_denominator = stored.float() if mode == 'stored_scale' else denominator |
| safe = torch.where(actual_denominator > 0, actual_denominator, torch.ones_like(actual_denominator)) |
| z = values/safe[..., None] |
| q = half_away_cpu(z).clamp(-bound, bound).int() |
| q = torch.where(actual_denominator[..., None] > 0, q, torch.zeros_like(q)) |
| integers.append(q); scales.append(stored) |
| return {'lo': pack_cpu(integers[0]), 'hi': pack_cpu(integers[0] >> 4), |
| 'q4': pack_cpu(integers[1]), 's8': scales[0], 's4': scales[1]} |
|
|
|
|
| def oracle(inputs, table, policy): |
| _, mode, clip, diagnostic = policy |
| x, rawdt, A, B, C, D, bias = [v.detach().cpu().float() for v in inputs] |
| table = table.cpu().long() |
| batch, length, heads, dim = x.shape |
| groups = B.shape[2] |
| group = torch.arange(heads)//(heads//groups) |
| indices = table[group][None].expand(batch, heads, 128) |
| previous = torch.zeros((batch, heads, dim, 128), dtype=torch.float32) |
| packed = quantize_cpu(previous, mode, clip) |
| outputs, caches, histories = [], [], [] |
| for t in range(length): |
| |
| if diagnostic is None: previous = decode_cpu(packed) |
| rdt = rawdt[:, t]+bias |
| delta = torch.where(rdt <= 20., torch.log(torch.exp(rdt)+1.), rdt) |
| decay = torch.exp(A[None]*delta) |
| b = B[:, t][:, group].gather(-1, indices) |
| c = C[:, t][:, group].gather(-1, indices) |
| updated = previous*decay[:, :, None, None] + (b*delta[..., None])[:, :, None]*x[:, t, :, :, None] |
| |
| y = (updated[..., :16]*c[:, :, None, :16]).sum(-1) |
| y = y + (updated[..., 16:80]*c[:, :, None, 16:80]).sum(-1) |
| y = y + ((updated[..., 80:]*c[:, :, None, 80:]).sum(-1)+x[:, t]*D[None, :, None]) |
| outputs.append(y.half()) |
| if diagnostic == 'prune_only': |
| previous = torch.cat((updated[..., :80].half().float(), torch.zeros_like(updated[..., 80:])), -1) |
| else: |
| packed = quantize_cpu(updated, mode, clip) |
| previous = decode_cpu(packed) |
| if diagnostic == 'quant_only': previous[..., 80:] = updated[..., 80:].half().float() |
| caches.append({k: v.clone() for k, v in packed.items()} if diagnostic is None else {'dense': previous.clone()}) |
| histories.append(previous.clone()) |
| return {'output': torch.stack(outputs, 1), 'caches': caches, 'decoded': torch.stack(histories)} |
|
|
|
|
| def controlled(length=33, special='ties'): |
| batch, heads, dim, groups = 2, 6, 5, 3 |
| table = torch.stack([torch.roll(torch.arange(127, -1, -1), 9*g) for g in range(groups)]).byte() |
| x = torch.empty((batch, length, heads, dim), dtype=torch.float16) |
| dt = torch.zeros((batch, length, heads), dtype=torch.float16) |
| A = torch.zeros(heads, dtype=torch.float32) |
| bias = torch.full((heads,), 32., dtype=torch.float32) |
| D = torch.tensor([(-1.)**h/16 for h in range(heads)], dtype=torch.float16) |
| B = torch.empty((batch, length, groups, 128), dtype=torch.float16); C = torch.empty_like(B) |
| n = torch.arange(128) |
| first = torch.tensor([-127., -126.5, -64., -2., -1.5, -.5, 0., .5, 1.5, 2., 63., 64., 126., 126.5, 127., 0.]) |
| second = torch.tensor([-7., -6.5, -4., -1.5, -.5, 0., .5, 1.5, 4., 6.5, 7.]).repeat(6)[:64] |
| wanted = torch.cat((first, second, ((torch.arange(48)%5)-2).float())) |
| if special == 'negative': wanted = -wanted.abs() |
| if special == 'zero': wanted.zero_() |
| if special == 'rounded_scales': wanted[0] = 127.0625; wanted[16] = 7.00390625 |
| if special == 'underflow': wanted = wanted.sign()*2**-19 |
| if special == 'subnormal': wanted = wanted.sign()*2**-11 |
| if special == 'mixed_scales': wanted[:16] = wanted[:16].sign()*2**-19; wanted[16:] = wanted[16:].sign()*2**-15 |
| if special == 'large_finite': wanted = wanted.sign()*131072.; wanted[80:] = 0. |
| for t in range(length): |
| sign = (1., 1., -1., -1.)[t%4] |
| for b in range(batch): |
| for h in range(heads): |
| for p in range(dim): |
| x[b, t, h, p] = (-1.)**(b+h+p)*2.**((t//4+b+h+p)%3-6) |
| for g in range(groups): |
| B[b, t, g, table[g].long()] = (wanted*sign/32).half() |
| C[b, t, g, table[g].long()] = (((n*3+t+b+g)%5-2).float()/256).half() |
| if special in ('underflow', 'mixed_scales'): x.fill_(2**-5) |
| if special == 'subnormal': x.fill_(2**-5) |
| if special == 'large_finite': x.fill_(1.); C.zero_(); D.zero_() |
| return [x, dt, A, B, C, D, bias], table |
|
|
|
|
| def random_inputs(geometry): |
| batch, length, heads, dim, groups = geometry |
| generator = torch.Generator().manual_seed(2026092906) |
| values = [(torch.randn(batch, length, heads, dim, generator=generator)*.08).half(), |
| (torch.randn(batch, length, heads, generator=generator)*.2-2).half(), |
| -torch.rand(heads, generator=generator).float()-.1, |
| (torch.randn(batch, length, groups, 128, generator=generator)*.2).half(), |
| (torch.randn(batch, length, groups, 128, generator=generator)*.2).half(), |
| (torch.randn(heads, generator=generator)*.1).half(), torch.zeros(heads, dtype=torch.float16)] |
| table = torch.stack([torch.randperm(128, generator=generator) for _ in range(groups)]).byte() |
| return values, table |
|
|
|
|
| def run(inputs, table, policy, pieces, *, capture=False): |
| from state_ppl_codec_v6 import allocate_state, scan, validate_finite_result, decode_state |
| _, mode, clip, diagnostic = policy |
| batch, length, heads, dim = inputs[0].shape |
| if sum(pieces) != length: raise ValueError('Partition differs from sequence') |
| state = allocate_state(batch, heads, dim, 128, 'cuda', diagnostic=diagnostic) |
| output, caches, decoded = [], [], []; start = 0 |
| for count in pieces: |
| part = [v[:, start:start+count] if i in (0, 1, 3, 4) else v for i, v in enumerate(inputs)] |
| y, unused = scan(*part, state, table, scale_mode=mode, int4_clip=clip, diagnostic=diagnostic) |
| if unused is not None: raise AssertionError('Unexpected persistent/statistics output') |
| validate_finite_result(state, y); output.append(y.cpu()) |
| if capture: |
| caches.append({k: v.cpu().clone() for k, v in state.tensors.items()}) |
| decoded.append(decode_state(state).cpu()) |
| start += count |
| result = {'output': torch.cat(output, 1), 'final': {k: v.cpu() for k, v in state.tensors.items()}, 'bytes': state.nbytes} |
| if capture: result.update(caches=caches, decoded=torch.stack(decoded)) |
| return result |
|
|
|
|
| def same_cache(a, b): return set(a) == set(b) and all(torch.equal(a[k], b[k]) for k in a) |
|
|
|
|
| def error(a, b): |
| difference = a.float()-b.float() |
| return dict(exact=torch.equal(a, b), different_elements=int((a != b).sum()), |
| max_absolute_error=float(difference.abs().max()), |
| relative_l2=float(torch.linalg.vector_norm(difference)/(torch.linalg.vector_norm(b.float())+1e-30))) |
|
|
|
|
|
|
| def scalar_arithmetic(kernel, ptx): |
| """Canonical emitted scalar arithmetic, excluding input addressing/layout.""" |
| lines, start = inspect.getsourcelines(kernel.fn) |
| raw = next(start+i for i, line in enumerate(lines) if 'dt = tl.load(DT' in line) |
| soft = next(start+i for i, line in enumerate(lines) if 'dt = tl.where(dt <=' in line) |
| decay = next(start+i for i, line in enumerate(lines) if 'decay = tl.exp(' in line) |
| selected, active, raw_active = [], False, False |
| for line in ptx.splitlines(): |
| match = re.search(r'\.loc\s+1\s+(\d+)\s+(\d+)', line) |
| if match: |
| number, column = map(int, match.groups()) |
| |
| |
| active = number in (soft, decay) or (number == 0 and active) |
| raw_active = number == raw |
| instruction = line.strip() |
| if not instruction or instruction.startswith(('.', '//', '$')): continue |
| if active or (raw_active and instruction.startswith('add.f32')): |
| selected.append(instruction) |
| names = {} |
| def register(match): |
| name = match.group() |
| if name not in names: names[name] = '%v'+str(len(names)) |
| return names[name] |
| canonical = re.sub(r'%(?:rs|rd|r|p)\d+', register, '\n'.join(selected)) |
| canonical = re.sub(r'\s+', ' ', canonical).strip() |
| return canonical, selected |
|
|
|
|
| def capture_scalar_values(inputs, table, artifact_dir): |
| """Capture actual GPU scalars and bind their PTX to both recurrence kernels.""" |
| from state_ppl_codec_v6 import _variant_scan, allocate_state |
| from mamba2_recall import state_codec as old |
|
|
| @triton.jit |
| def scalar_probe(DT, A, DB, Delta, Decay, H: tl.constexpr): |
| i = tl.program_id(0) |
| h = i % H |
| av = tl.load(A+h).to(tl.float32) |
| bias = tl.load(DB+h).to(tl.float32) |
| dt = tl.load(DT+i).to(tl.float32) + bias |
| dt = tl.where(dt <= 20., tl.math.log(tl.math.exp(dt)+1.), dt) |
| decay = tl.exp(av*dt) |
| tl.store(Delta+i, dt) |
| tl.store(Decay+i, decay) |
|
|
| x, dt, A, B, C, D, bias = [v.cuda().contiguous() for v in inputs] |
| permutation = table.cuda().contiguous() |
| batch, length, heads, dim = x.shape |
| groups = B.shape[2] |
| delta, decay = torch.empty_like(dt, dtype=torch.float32), torch.empty_like(dt, dtype=torch.float32) |
| probe = scalar_probe[(dt.numel(),)](dt, A, bias, delta, decay, heads, num_warps=4) |
| |
| |
| one = [v[:, :1].contiguous() if i in (0, 1, 3, 4) else v for i, v in enumerate((x, dt, A, B, C, D, bias))] |
| packed = old.allocate_state(batch, heads, dim, 128, 'sq3p25', 'cuda') |
| dense = allocate_state(batch, heads, dim, 128, 'cuda', diagnostic='legacy_emulation') |
| y, z = torch.empty_like(one[0]), torch.empty_like(one[0]) |
| launch = (batch, heads, triton.cdiv(dim, 16)) |
| original = old._sq_scan[launch](*one, permutation, *(packed.tensors[k] for k in ('lo', 'hi', 'q4', 's8', 's4')), |
| y, 1, heads, dim, groups, 16, num_warps=4) |
| backing = dense.tensors['dense'] |
| emulated = _variant_scan[launch](*one, permutation, *([backing]*5), backing, z, |
| 1, heads, dim, groups, 16, False, 1., 3, num_warps=4) |
| artifacts, canonical = {}, {} |
| for name, kernel, compiled in (('packed', old._sq_scan, original), ('dense', _variant_scan, emulated), ('scalar_probe', scalar_probe, probe)): |
| ptx = compiled.asm['ptx'] |
| path = artifact_dir/(name+'.ptx'); path.write_text(ptx) |
| text, instructions = scalar_arithmetic(kernel, ptx) |
| canonical[name] = text |
| artifacts[name] = {'file': str(path), 'sha256': sha(path), 'instructions': len(instructions), |
| 'scalar_sha256': hashlib.sha256(text.encode()).hexdigest()} |
| equal = len(set(canonical.values())) == 1 |
| if not equal: raise AssertionError('Captured scalar PTX differs from actual packed/dense recurrence') |
| if not bool(torch.isfinite(delta).all() and torch.isfinite(decay).all()): |
| raise AssertionError('Scalar error-bound scope requires finite delta/decay') |
| return delta.cpu(), decay.cpu(), {'exact': equal, 'artifacts': artifacts} |
|
|
|
|
| def dense_emulation_bound(inputs, table, packed, dense, artifact_dir): |
| """Independent FP64 gamma bound; no tolerance fitted to measured drift. |
| |
| Per-path recurrence(3)+C product(1)+sum(N-1)+outeradds(2) |
| gives gamma21/69/69. The xD path has at most3 roundings. Underflow |
| floor covers fewer than1024 operations, <=3 flush events each, and |
| downstream gains from x and C. Both output-rounding cells are included. |
| """ |
| artifact_dir.mkdir(parents=True, exist_ok=False) |
| delta, decay, scalar_proof = capture_scalar_values(inputs, table, artifact_dir) |
| x, rawdt, A, B, C, D, bias = [v.double() for v in inputs] |
| batch, length, heads, dim = x.shape |
| group = torch.arange(heads)//(heads//B.shape[2]) |
| indices = table.long()[group][None].expand(batch, heads, 128) |
| u32, u64 = 2.**-24, 2.**-53 |
| gamma = lambda n: n*u32/(1.-n*u32) |
| upward = 1.+1024*u64/(1.-1024*u64) |
| bounds = [] |
| for t in range(length): |
| previous = torch.zeros_like(packed['decoded'][0], dtype=torch.float64) if t == 0 else packed['decoded'][t-1].double() |
| b = B[:, t][:, group].gather(-1, indices) |
| c = C[:, t][:, group].gather(-1, indices) |
| old = previous.abs()*decay[:, t, :, None, None].double().abs() |
| drive = (b*delta[:, t, :, None].double()).abs()[:, :, None]*x[:, t, :, :, None].abs() |
| terms = (old+drive)*c[:, :, None].abs() |
| e = gamma(21)*terms[..., :16].sum(-1)+gamma(69)*terms[..., 16:80].sum(-1)+gamma(69)*terms[..., 80:].sum(-1) |
| e += gamma(3)*(x[:, t]*D[None, :, None]).abs() |
| e += 4096.*2.**-126*(1.+x[:, t].abs())*(1.+c.abs().sum(-1)[:, :, None]) |
| bounds.append(e) |
| e = torch.stack(bounds, 1) |
| a, b = packed['output'].double(), dense['output'].double() |
| def radius(value): |
| _, exponent = torch.frexp(value.abs()) |
| return torch.pow(2., exponent.double()-12.).clamp_min(2.**-25) |
| |
| ra = torch.where(a == 0, torch.full_like(a, 2.**-25), radius(a)) |
| rb = torch.where(b == 0, torch.full_like(b, 2.**-25), radius(b)) |
| bound = torch.nextafter((2.*e+ra+rb)*upward, torch.full_like(e, float('inf'))) |
| difference = (a-b).abs() |
| finite = bool(torch.isfinite(a).all() and torch.isfinite(b).all() and torch.isfinite(bound).all()) |
| carry_exact = torch.equal(packed['decoded'], dense['decoded']) |
| failures = int((difference > bound).sum()) |
| path = artifact_dir/'evidence.pt' |
| torch.save({'inputs': inputs, 'table': table, 'dt_actual': delta, 'decay_actual': decay, |
| 'packed_output': packed['output'], 'dense_output': dense['output'], |
| 'packed_carry': packed['decoded'], 'dense_carry': dense['decoded'], |
| 'packed_caches': packed['caches'], 'dense_caches': dense['caches'], 'bound': bound}, path) |
| facts = {'all_token_carries_exact': carry_exact, 'finite_scope': finite, 'bound_failures': failures, |
| 'max_error_to_bound_ratio': float((difference/bound).max()), |
| 'max_bound': float(bound.max()), 'output_error': error(packed['output'], dense['output']), |
| 'scalar_ptx_equivalence': scalar_proof, 'evidence': {'file': str(path), 'sha256': sha(path), 'bytes': path.stat().st_size}, |
| 'bound_formula': '2*(gamma21*T8+gamma69*T4+gamma69*Td+gamma3*abs(xD)+4096*2^-126*(1+abs(x))*(1+sum(abs(C))))+r16(a)+r16(b)', |
| 'float64_upward_guard': 'multiply by1+gamma1024(u64), then nextafter(+inf)'} |
| return carry_exact and finite and failures == 0 and scalar_proof['exact'], facts |
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--out', type=Path, required=True) |
| parser.add_argument('--cpu-only', action='store_true') |
| args = parser.parse_args() |
| if args.out.exists(): raise FileExistsError(args.out) |
| args.out.parent.mkdir(parents=True, exist_ok=True) |
| report = {'format': 'MAMBA2_STATE_PPL_CODEC_CPU_PREPARATION_V1' if args.cpu_only else 'MAMBA2_STATE_PPL_CODEC_CHECK_V1', |
| 'complete': False, 'passed': False, 'mode': 'cpu_preparation' if args.cpu_only else 'gpu', |
| 'protocol_sha256': PROTOCOL_SHA, 'started_unix': time.time(), 'checks': [], |
| 'code_sha256': {p: sha(ROOT/p) for p in ('scripts/state_ppl_codec_v6.py', |
| 'scripts/check_state_ppl_codec_v6.py', 'mamba2_recall/state_codec.py', 'mamba2_recall/state_quant.py')}, |
| 'scope': 'Independent controlled arithmetic/packing plus GPU kernel partition and legacy replay; no model PPL/MK measurements', |
| 'criteria': {'controlled': 'exact every-token output/packed bytes/FP16 scales/decoded carry', |
| 'random_partition': 'exact full/segmented/tokenwise outputs and final cache', |
| 'random_one_token_cpu_output_relative_l2_bound': .002, |
| 'stored_scale_division': 'Correctly rounded FP32 quotient before half-away rounding; zero-scale integer0', |
| 'stored_scale_denominator': 'FP32 clipping multiplication then correctly rounded FP32 division before FP16 round-to-even storage', |
| 'diagnostic_legacy_emulation': 'controlled output exact; random every-token carry exact and readout within recurrence-inclusive gamma21/69/69/3 plus FP16-cell and FP32-underflow bound'}} |
| def check(name, condition, **facts): |
| row = dict(name=name, **{'pass': bool(condition)}, **facts); report['checks'].append(row) |
| print(json.dumps(row), flush=True) |
| if not condition: raise AssertionError(name) |
| try: |
| torch.set_num_threads(4) |
| check('frozen_protocol', sha(ROOT/'docs/STATE_PPL_V6_PROTOCOL.md') == PROTOCOL_SHA) |
| halves = torch.tensor([.5, 3.5], dtype=torch.float32) |
| below = torch.nextafter(halves, torch.zeros_like(halves)) |
| boundary_values = torch.stack((halves, below, -halves, -below), -1).flatten() |
| boundary_expected = torch.tensor([1, 0, -1, 0, 4, 3, -4, -3], dtype=torch.int32) |
| check('CPU-half-away-near-boundary', torch.equal(half_away_cpu(boundary_values).int(), boundary_expected)) |
| if not args.cpu_only: |
| if not torch.cuda.is_available(): raise RuntimeError('CUDA required for actual codec checks') |
| torch.backends.cuda.matmul.allow_tf32 = False; torch.backends.cudnn.allow_tf32 = False |
| report['gpu'] = torch.cuda.get_device_name() |
| from state_ppl_codec_v6 import quantize_stored_probe |
| |
| |
| scale = torch.tensor(.037933349609375, dtype=torch.float16) |
| exact_value = scale.float()*3.5 |
| below_value = torch.nextafter(exact_value, torch.tensor(0.)) |
| probe_values = torch.cat((boundary_values, torch.stack((exact_value, below_value, -exact_value, -below_value)), torch.tensor([1., -1., 0.]))) |
| probe_scales = torch.cat((torch.ones(8, dtype=torch.float16), scale.repeat(4), torch.zeros(3, dtype=torch.float16))) |
| probe_expected = torch.cat((boundary_expected, torch.tensor([4, 3, -4, -3, 0, 0, 0], dtype=torch.int32))) |
| for limit in (7, 127): |
| result = quantize_stored_probe(probe_values.cuda(), probe_scales.cuda(), limit).cpu() |
| check('GPU-stored-scale-half-boundaries-'+str(limit), torch.equal(result, probe_expected), |
| actual_codes=result.tolist(), expected_codes=probe_expected.tolist()) |
| |
| for special in ('ties', 'negative', 'zero', 'rounded_scales', 'underflow', 'subnormal', 'mixed_scales', 'large_finite'): |
| inputs, table = controlled(1, special) |
| for policy in POLICIES: |
| if special == 'large_finite' and policy[3] == 'prune_only': continue |
| expected = oracle(inputs, table, policy) |
| check('CPU-fixture-finite-'+special+'-'+policy[0], bool(torch.isfinite(expected['output']).all() |
| and torch.isfinite(expected['decoded']).all())) |
| if special == 'underflow' and policy[1] == 'stored_scale': |
| buffers = expected['caches'][0] |
| check('CPU-zero-scale-integer-zero-'+policy[0], |
| all(bool((v == 0).all()) for v in buffers.values())) |
| if special == 'subnormal' and policy[3] is None: |
| buffers = expected['caches'][0] |
| check('CPU-positive-subnormal-scale-'+policy[0], |
| any(bool(((buffers[k] > 0) & (buffers[k] < torch.finfo(torch.float16).tiny)).any()) for k in ('s8', 's4'))) |
| if args.cpu_only: continue |
| gpu_inputs, gpu_table = [v.cuda() for v in inputs], table.cuda() |
| actual = run(gpu_inputs, gpu_table, policy, [1], capture=True) |
| check('GPU-independent-exact-'+special+'-'+policy[0], |
| torch.equal(actual['output'], expected['output']) |
| and same_cache(actual['caches'][0], expected['caches'][0]) |
| and torch.equal(actual['decoded'], expected['decoded']), |
| output_error=error(actual['output'], expected['output']), carry_error=error(actual['decoded'], expected['decoded'])) |
| inputs, table = controlled(65) |
| controlled_gpu = {} |
| for policy in POLICIES: |
| expected = oracle(inputs, table, policy) |
| check('CPU-multistep-finite-'+policy[0], bool(torch.isfinite(expected['output']).all() and torch.isfinite(expected['decoded']).all())) |
| if args.cpu_only: continue |
| actual = run([v.cuda() for v in inputs], table.cuda(), policy, [1]*65, capture=True) |
| controlled_gpu[policy[0]] = actual |
| check('GPU-65token-independent-exact-'+policy[0], torch.equal(actual['output'], expected['output']) |
| and all(same_cache(a, b) for a, b in zip(actual['caches'], expected['caches'])) |
| and torch.equal(actual['decoded'], expected['decoded']), |
| output_error=error(actual['output'], expected['output']), carry_error=error(actual['decoded'], expected['decoded'])) |
| if not args.cpu_only: |
| from state_ppl_codec_v6 import allocate_state, scan, validate_options, validate_finite_result |
| from mamba2_recall import state_codec as legacy |
| for geometry in ((2, 17, 8, 19, 2), (1, 9, 6, 3, 3), (2, 5, 4, 1, 1)): |
| inputs, table = random_inputs(geometry); b, length, h, p, g = geometry |
| gpu_inputs, gpu_table = [v.cuda() for v in inputs], table.cuda() |
| results = {} |
| for policy in POLICIES: |
| full = run(gpu_inputs, gpu_table, policy, [length]) |
| parts = run(gpu_inputs, gpu_table, policy, [1, 3, length-4]) |
| tokens = run(gpu_inputs, gpu_table, policy, [1]*length, capture=True) |
| results[policy[0]] = tokens |
| check('partition-'+str(geometry)+'-'+policy[0], torch.equal(full['output'], parts['output']) |
| and torch.equal(full['output'], tokens['output']) and same_cache(full['final'], parts['final']) |
| and same_cache(full['final'], tokens['final'])) |
| row_bytes = 52 if policy[3] is None else 512 |
| check('actual-bytes-'+str(geometry)+'-'+policy[0], full['bytes'] == b*h*p*row_bytes, |
| row_bytes=row_bytes, state_bytes=full['bytes'], is_3p25_candidate=policy[3] is None) |
| single = [v[:, :1] if i in (0, 1, 3, 4) else v for i, v in enumerate(inputs)] |
| expected = oracle(single, table, policy) |
| e = error(tokens['output'][:, :1], expected['output']) |
| check('random-one-token-oracle-'+str(geometry)+'-'+policy[0], e['relative_l2'] <= .002, **e) |
| bounded, bound_facts = dense_emulation_bound(inputs, table, results['legacy'], results['legacy_emulation'], |
| args.out.parent/(args.out.stem+'_details')/('geometry_'+'_'.join(map(str, geometry)))) |
| check('dense-legacy-emulation-'+str(geometry), bounded, **bound_facts) |
| original_state = legacy.allocate_state(b, h, p, 128, 'sq3p25', 'cuda') |
| original_output, _ = legacy.scan(*gpu_inputs, original_state, gpu_table) |
| check('delegated-legacy-'+str(geometry), torch.equal(original_output.cpu(), results['legacy']['output']) |
| and same_cache({k: v.cpu() for k, v in original_state.tensors.items()}, results['legacy']['final'])) |
| first = results['stored_scale']['final'] |
| check('INT8-unchanged-by-INT4-clipping-'+str(geometry), all( |
| all(torch.equal(first[k], results[name]['final'][k]) for k in ('lo', 'hi', 's8')) |
| for name in ('clip4_095', 'clip4_090', 'clip4_080'))) |
| for diagnostic in (None, 'prune_only', 'quant_only', 'legacy_emulation'): |
| state = allocate_state(1, 128, 64, 128, 'cuda', diagnostic=diagnostic) |
| total = 56*state.nbytes+56*10240*4*2+56*8*128 |
| check('production-budget-'+str(diagnostic), total == (28499968 if diagnostic is None else 239525888), |
| actual_state_bytes=56*state.nbytes, table_bytes=57344, total_bytes=total, |
| is_3p25_candidate=diagnostic is None) |
| |
| inputs, table = controlled(1) |
| inputs[0].fill_(65504.); inputs[3].fill_(65504.); inputs[4].zero_(); inputs[5].zero_() |
| for policy in POLICIES: |
| state = allocate_state(2, 6, 5, 128, 'cuda', diagnostic=policy[3]) |
| output, _ = scan(*[v.cuda() for v in inputs], state, table.cuda(), |
| scale_mode=policy[1], int4_clip=policy[2], diagnostic=policy[3]) |
| rejected = False |
| try: validate_finite_result(state, output) |
| except FloatingPointError: rejected = True |
| check('reject-nonfinite-carry-or-scale-'+policy[0], rejected) |
| rejected = False |
| try: validate_options('legacy', .95) |
| except ValueError: rejected = True |
| check('reject-unfrozen-legacy-clipping-policy', rejected) |
| report['cuda_initialized'] = torch.cuda.is_initialized() |
| if args.cpu_only: check('CUDA-remained-uninitialized', not report['cuda_initialized']) |
| report.update(complete=True, passed=True) |
| except BaseException as exc: |
| report['error'] = repr(exc) |
| raise |
| finally: |
| report['finished_unix'] = time.time() |
| args.out.write_text(json.dumps(report, indent=2, allow_nan=False)+'\n') |
|
|
|
|
| if __name__ == '__main__': main() |
|
|
