scripts/check_state_ppl_codec_v6.py

#!/usr/bin/env python3
"""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):
    # Avoid floor(abs(z)+.5): that addition can round nextafter(.5,0) up to1.
    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):
        # Packed oracle reconstructs the ACTUALLY STORED bytes/scales each token.
        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]
        # Three separate reductions independently represent the codec contract.
        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())
            # Nested JIT functions are dedented; the libdevice fallback's
            # source column changes although its emitted instructions do not.
            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-token kernels are the actual path already checked against full/chunk
    # executions. Capture their emitted scalar code without changing their math.
    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)
    # frexp(0) returns exponent0, so explicitly set its subnormal cell radius.
    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
            # Include the exact non-power-of-two FP16 denominator that exposed
            # approximate reciprocal division in preserved failed attempt1.
            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())
        # No implementation codec functions are used to construct these expected results.
        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)
            # Finite inputs can overflow the stored scale/carry; reject them explicitly.
            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()