exp: characterize KONI numerical drift and output-head precision limits

opt_prime/examples/pd/http_matrix/.research/events.jsonl

@@ -5,3 +5,4 @@
 {"id": "328b8349e9974a059ddc7e1a8a485e2c", "ts": "2026-10-23T21:32:23+09:00", "task": "Tokenizer EOS in batched serving (aicomp #229)", "result": "Real ATOM KONI HTTP also passes5/5 baseline/fixed/ignore/stream/reuse checks using same production factory: config-only length vs union stop. GPU+ATOM now10/10 expected outcomes; no SDK/artifact rebuild.", "next": "Update PR230 with NPU evidence; finish KONI medium-context RNGD handoff and broader PD work.", "state": "active", "done": false, "blocker": "", "metrics": {}, "params": {"backend": "atom", "kv_blocks": "1"}, "artifacts": [{"path": "analysis/serving-matrix-201/eos229-koni-atom-http.json", "size": 5199, "sha256": "4735255db379ff00"}], "experiments": ["experiments/2026-10-23-serving-eos229.md"], "labels": ["bug"], "review": true, "todos": [], "note": "", "agent": "codex", "host": "xeon-max", "branch": "main", "commit": "351d39d-dirty"}
 {"id": "1020452aea8445f080108c1ed7f6d643", "ts": "2026-10-24T04:35:37+09:00", "task": "Dense hybrid GPU KV import (aicomp #258)", "result": "Dense hybrid receiver fix:31focused and505full CPU tests passed,34skipped,253subtests,151.69s. Preparing isolated PR and fresh Gemma live rerun; hardware and numerical acceptance not yet claimed.", "next": "Run corrected Gemma RNGD/GPU matrix preserving original failure; complete lifecycle and numerical checks", "state": "active", "done": false, "blocker": "", "metrics": {}, "params": {}, "artifacts": [], "experiments": ["experiments/2026-10-24-dense-hybrid258.md"], "labels": ["bug"], "review": false, "todos": [], "note": "", "agent": "codex", "host": "xeon-max", "branch": "main", "commit": "351d39d-dirty"}
 {"id": "551af4e77bca4e67bf0512304142ac48", "ts": "2026-10-24T04:41:11+09:00", "task": "Dense hybrid GPU KV import (aicomp #258)", "result": "Corrected live Gemma completed24batches;20/20token checks and60/60timed texts exact to nativeGPU,84remote/0local/0errors. Shape crash fixed on hardware. Qwen0.6 GPU lifecycle passed; Qwen4 next. Owned ATOM Pod confirmed absent.", "next": "Publish hardware evidence PR259; finish remaining GPU fault coverage and numerical boundary work", "state": "active", "done": false, "blocker": "", "metrics": {}, "params": {}, "artifacts": [{"path": "analysis/serving-matrix-201/local258-gemma-summary.json", "size": 1288, "sha256": "3afbcc7e61bb3239"}], "experiments": ["experiments/2026-10-24-dense-hybrid258.md"], "labels": ["bug"], "review": false, "todos": [], "note": "", "agent": "codex", "host": "xeon-max", "branch": "main", "commit": "351d39d-dirty"}
+{"id": "232c1a48c9bf49178446feda8f64b935", "ts": "2026-10-24T05:09:31+09:00", "task": "KONI numerical compatibility (aicomp #261)", "result": "Unforced head-onlyFP32 control completed with baseline reproduction: KONIshort firstdiff14 vs original64; Qwen4long firstdiff13 vs107; KONIlong matchedprefill remains exact. Reject head-onlyprecision as generalfix; production unchanged.", "next": "Publish complete numerical controls; align compatibility criteria and evaluate held-out quality instead of fixture-specific tie fixes", "state": "active", "done": false, "blocker": "", "metrics": {}, "params": {}, "artifacts": [{"path": "analysis/serving-matrix-201/head-fp32-free-generation.json", "size": 29033, "sha256": "83906dd31756d2a9"}], "experiments": ["experiments/2026-10-24-koni-numerics261.md"], "labels": ["analysis"], "review": false, "todos": [], "note": "", "agent": "codex", "host": "xeon-max", "branch": "main", "commit": "351d39d-dirty"}

opt_prime/examples/pd/http_matrix/STATUS.md

@@ -41,3 +41,7 @@ Candidate exports airm_pd_producer_seconds_total{stage="prefill"|"transfer"} fro
 #258 dense hybrid KV: actual Gemma failure reproduced; mapping/group validation corrected.31focused and505full CPU tests pass; live rerun pending.
 
 #258 corrected live Gemma32ce5d5:24batches,20/20token checks,60/60timed texts exact;84remote/0local/0errors. Shape crash fixed in recorded hardware matrix; overall numerical/quality boundaries remain.
+
+#261 KONI numerical evidence: original live drift beforeEOS, shortKV exact repeat/replay, padded-shape control, forced-prefix BF16/FP32 and head-onlyFP32 probes. No production fix/general equivalence claim.
+
+Free-generation control rejects head-onlyFP32 as generalfix: KONIshort firstdiff14 andQwen4long13; no production option added.

opt_prime/examples/pd/http_matrix/evidence261/check_head_precision_free.py

@@ -0,0 +1,66 @@
+"""Check unforced128-token generation after changing only output projection precision."""
+import gc
+import json
+from pathlib import Path
+from concurrent.futures import Future
+import torch
+from transformers import AutoModelForCausalLM
+from opt_prime.accel.pd import PDRequest
+from opt_prime.accel.pd.batching import ContinuousBatcher
+from opt_prime.accel.pd.pd_batching import PDBatcher, PrefillOut
+root=Path(__file__).resolve().parent
+output=root/'head-fp32-free-generation.json'
+if output.exists(): raise FileExistsError(output)
+torch.set_num_threads(4)
+torch.backends.cuda.matmul.allow_tf32=False
+rows=[]
+for label, model_path, cases in [
+    ('koni','/tmp/serving-matrix-models/KONI-8B',[(128,256,'koni261-short-kv.pt'),(8192,8192,'rngd-koni-8b-8192-parts221-layout.pt')]),
+    ('qwen4b','/tmp/serving-matrix-models/Qwen3-4B',[(8192,8192,'rngd-qwen3-4b-8192.pt')])]:
+    data=json.loads((root/f'local201-{label}-matrix-a.json').read_text())
+    model=AutoModelForCausalLM.from_pretrained(model_path,dtype=torch.bfloat16,attn_implementation='sdpa',local_files_only=True).cuda().eval()
+    head=model.get_output_embeddings()
+    original_forward=head.forward
+    weight_fp32=head.weight.detach().float()
+    bias_fp32=None if head.bias is None else head.bias.detach().float()
+    calls=[0]
+    def project(hidden):
+        calls[0]+=1
+        return torch.nn.functional.linear(hidden.float(),weight_fp32,bias_fp32)
+    for n,chunk,name in cases:
+        export=torch.load(root/name,weights_only=True,mmap=True)
+        assert export['ids']==data['prompts'][str(n)]
+        live=next(b['rows'][0]['tokens'] for b in data['correctness_batches'] if b['input_tokens']==n and b['output_tokens']==128)
+        if label=='qwen4b':
+            expected=json.loads((root/'local201-qwen4b-native-full.json').read_text())['references']['8192:128']['tokens']
+        elif n==8192:
+            expected=next(r['tokens'] for r in json.loads((root/'local201-koni-diagnostic.json').read_text()) if r['mode']=='native8192')
+        else: expected=data['references']['128:128']['tokens']
+        class Replay:
+            def submit(self,request,send_from=0,slot=None):
+                assert list(request.input_ids)==export['ids'] and send_from==0
+                f=Future()
+                token=export['token'] if 'token' in export else int(export['logits'].argmax())
+                f.set_result(PrefillOut(token,0,export['keys'],export['values']))
+                return f
+        result=dict(model=label,input_tokens=n,native_chunk=chunk,rows=[])
+        for precision in ['bf16','fp32_head']:
+            head.forward=original_forward if precision=='bf16' else project
+            for mode in ['native','retained_rngd_kv']:
+                kw=dict(slots=4,max_context=9216,kv='paged',page_size=64,prefix_cache=False)
+                b=ContinuousBatcher(model,prefill_chunk=chunk,**kw) if mode=='native' else PDBatcher(model,Replay(),**kw)
+                f=b.submit(PDRequest(mode,tuple(export['ids']),128));b.run_until_idle();tokens=f.result()
+                if precision=='bf16': assert tokens==(expected if mode=='native' else live), 'baseline reproduction failed'
+                else: assert calls[0]>0 and next(model.parameters()).dtype==torch.bfloat16
+                result['rows'].append(dict(precision=precision,mode=mode,tokens=tokens))
+                del b,f
+                gc.collect();torch.cuda.empty_cache()
+        a,b=[r['tokens'] for r in result['rows'] if r['precision']=='fp32_head']
+        result.update(head_fp32_equal=a==b,first_difference=next((i for i,(x,y) in enumerate(zip(a,b)) if x!=y),None))
+        rows.append(result)
+        output.write_text(json.dumps(dict(complete=False,rows=rows),indent=2)+'\n')
+        print(json.dumps({k:v for k,v in result.items() if k!='rows'}),flush=True)
+    head.forward=original_forward
+    del model,head,original_forward,weight_fp32,bias_fp32
+    gc.collect();torch.cuda.empty_cache()
+output.write_text(json.dumps(dict(complete=True,rows=rows),indent=2)+'\n')

opt_prime/examples/pd/http_matrix/evidence261/diagnose_koni261_head.py

@@ -0,0 +1,112 @@
+"""Verify short-KV replay, then compare KONI logits under an identical64-token prefix."""
+import gc
+import json
+from pathlib import Path
+from concurrent.futures import Future
+import torch
+from transformers import AutoModelForCausalLM
+from opt_prime.accel.pd import PDRequest
+from opt_prime.accel.pd.batching import ContinuousBatcher
+from opt_prime.accel.pd.pd_batching import PDBatcher, PrefillOut
+
+root = Path(__file__).resolve().parent
+output = root/'koni261-head-fp32.json'
+if output.exists():
+    raise FileExistsError(output)
+export = torch.load(root/'koni261-short-kv.pt', weights_only=True, mmap=True)
+data = json.loads((root/'local201-koni-matrix-a.json').read_text())
+assert export['ids'] == data['prompts']['128']
+native = data['references']['128:128']['tokens']
+live = next(b['rows'][0]['tokens'] for b in data['correctness_batches']
+            if b['input_tokens']==128 and b['output_tokens']==128)
+forced = native[:64]
+assert forced == live[:64] and native[64] != live[64]
+torch.set_num_threads(4)
+torch.backends.cuda.matmul.allow_tf32 = False
+result = dict(complete=False, replay_verified=False, rows=[])
+
+class Replay:
+    def submit(self, request, send_from=0, slot=None):
+        assert list(request.input_ids) == export['ids'] and send_from == 0
+        future = Future()
+        future.set_result(PrefillOut(export['token'], 0, export['keys'], export['values']))
+        return future
+
+def make(model, mode):
+    kw = dict(slots=4, max_context=9216, kv='paged', page_size=64, prefix_cache=False)
+    return PDBatcher(model, Replay(), **kw) if mode == 'retained_rngd_kv' else ContinuousBatcher(model, prefill_chunk=256, **kw)
+
+for dtype in [torch.bfloat16]:
+    model = AutoModelForCausalLM.from_pretrained('/tmp/serving-matrix-models/KONI-8B',
+        dtype=dtype, attn_implementation='sdpa', local_files_only=True).cuda().eval()
+    if dtype == torch.bfloat16:
+        for mode, expected in [('native128',native),('retained_rngd_kv',live)]:
+            batcher = make(model, mode)
+            future = batcher.submit(PDRequest(mode, tuple(export['ids']),128))
+            batcher.run_until_idle()
+            assert future.result() == expected, 'Original output not reproduced: '+mode
+            del batcher, future
+            gc.collect()
+            torch.cuda.empty_cache()
+        result['replay_verified'] = True
+        output.write_text(json.dumps(result,indent=2)+'\n')
+    head = model.get_output_embeddings()
+    original_head_forward = head.forward
+    weight_fp32 = head.weight.detach().float()
+    bias_fp32 = None if head.bias is None else head.bias.detach().float()
+    head_calls = [0]
+    def fp32_projection(hidden):
+        head_calls[0] += 1
+        return torch.nn.functional.linear(hidden.float(), weight_fp32, bias_fp32)
+    head.forward = fp32_projection
+    assert next(model.parameters()).dtype == torch.bfloat16 and head.weight.dtype == torch.bfloat16
+    for mode in ['native128','retained_rngd_kv']:
+        batcher = make(model, mode)
+        emitted, decoded = batcher._emit, batcher._decoded
+        probes, kv_error = [], []
+
+        def emit(request, token):
+            index = len(request.tokens)
+            if index == 0 and mode == 'native128':
+                keys, values = batcher.kv.read_span(request.slot, 0, len(export['ids']))
+                for layer in range(keys.shape[0]):
+                    row = dict(layer=layer)
+                    for name, actual in [('keys',keys),('values',values)]:
+                        ref = actual[layer].float()
+                        diff = export[name][layer].to(ref.device).float() - ref
+                        row[name+'_relative_l2'] = float(torch.linalg.vector_norm(diff)/torch.linalg.vector_norm(ref).clamp_min(1e-20))
+                        row[name+'_max_abs'] = float(diff.abs().max())
+                    kv_error.append(row)
+            return emitted(request, forced[index] if index < 64 else token)
+
+        def decode(requests, lengths, logits):
+            for i, request in enumerate(requests):
+                index = len(request.tokens)
+                if index in [63,64]:
+                    assert request.tokens == forced[:index]
+                    values, tokens = logits[i].float().topk(8)
+                    probes.append(dict(output_index=index,finite=bool(torch.isfinite(logits[i]).all()),
+                        top_tokens=tokens.tolist(),top_logits=values.tolist(),native_branch_token=native[index],
+                        live_branch_token=live[index],native_branch_logit=float(logits[i,native[index]]),
+                        live_branch_logit=float(logits[i,live[index]])))
+            return decoded(requests,lengths,logits)
+
+        batcher._emit, batcher._decoded = emit, decode
+        future = batcher.submit(PDRequest(mode,tuple(export['ids']),65))
+        batcher.run_until_idle()
+        tokens = future.result()
+        assert tokens[:64] == forced and len(probes)==2 and all(p['finite'] for p in probes)
+        assert head_calls[0] > 0
+        result['rows'].append(dict(dtype=str(dtype),head_dtype='float32',head_calls=head_calls[0],mode=mode,selected_token65=tokens[-1],
+                                   probes=probes,kv_difference_from_native=kv_error))
+        output.write_text(json.dumps(result,indent=2)+'\n')
+        print(json.dumps(dict(dtype=str(dtype),mode=mode,token65=tokens[-1],probe=probes[-1])),flush=True)
+        batcher._emit, batcher._decoded = emitted, decoded
+        del batcher, emitted, decoded, future
+        gc.collect()
+        torch.cuda.empty_cache()
+    del model
+    gc.collect()
+    torch.cuda.empty_cache()
+result['complete'] = True
+output.write_text(json.dumps(result,indent=2)+'\n')

opt_prime/examples/pd/http_matrix/evidence261/diagnose_koni_short_padding.py

@@ -0,0 +1,60 @@
+"""Compare real-length and RNGD-shaped padded GPU prefill without changing production sampling."""
+import json
+from pathlib import Path
+from concurrent.futures import Future
+import torch
+from transformers import AutoModelForCausalLM
+from opt_prime.accel.pd import PDRequest
+from opt_prime.accel.pd.batching import ContinuousBatcher
+from opt_prime.accel.pd.pd_batching import PDBatcher, PrefillOut
+
+root = Path(__file__).resolve().parent
+output = root/'koni-short-padding-diagnostic.json'
+if output.exists():
+    raise FileExistsError(output)
+data = json.loads((root/'local201-koni-matrix-a.json').read_text())
+ids = data['prompts']['128']
+assert len(ids) == 128
+native = data['references']['128:128']['tokens']
+live = next(b['rows'][0]['tokens'] for b in data['correctness_batches']
+            if b['input_tokens'] == 128 and b['output_tokens'] == 128)
+torch.set_num_threads(4)
+model = AutoModelForCausalLM.from_pretrained('/tmp/serving-matrix-models/KONI-8B',
+    dtype=torch.bfloat16, attn_implementation='sdpa', local_files_only=True).cuda().eval()
+rows = []
+
+class PaddedGPU:
+    def __init__(self, width):
+        self.width = width
+
+    @torch.inference_mode()
+    def submit(self, request, send_from=0, slot=None):
+        assert list(request.input_ids) == ids and send_from == 0
+        # The research RNGD _DevicePrefillStage repeats the last token to its bucket width.
+        padded = ids + [ids[-1]] * (self.width - len(ids))
+        out = model(torch.tensor([padded], device='cuda:0'), use_cache=True,
+                    logits_to_keep=torch.tensor([len(ids)-1], device='cuda:0'))
+        keys = torch.stack([layer.keys[0, :, :len(ids)].contiguous() for layer in out.past_key_values.layers])
+        values = torch.stack([layer.values[0, :, :len(ids)].contiguous() for layer in out.past_key_values.layers])
+        result = PrefillOut(int(out.logits[0, 0].argmax()), 0, keys, values)
+        future = Future()
+        future.set_result(result)
+        return future
+
+for mode, width in [('native128', 128), ('hf128_import', 128), ('hf8192_padded_import', 8192)]:
+    kw = dict(slots=4, max_context=9216, kv='paged', page_size=64, prefix_cache=False)
+    batcher = ContinuousBatcher(model, prefill_chunk=256, **kw) if mode == 'native128' else PDBatcher(model, PaddedGPU(width), **kw)
+    future = batcher.submit(PDRequest(mode, tuple(ids), 128))
+    batcher.run_until_idle()
+    tokens = future.result()
+    assert len(tokens) == 128
+    if mode == 'native128':
+        assert tokens == native, 'native baseline changed; stop diagnostic'
+    rows.append(dict(mode=mode, physical_prefill_tokens=width, logical_prompt_tokens=128,
+                     tokens=tokens, equals_native=tokens == native, equals_live=tokens == live,
+                     first_diff_native=next((i for i,(a,b) in enumerate(zip(tokens,native)) if a != b), None)))
+    output.write_text(json.dumps(dict(complete=False, rows=rows), indent=2)+'\n')
+    print(json.dumps({k:v for k,v in rows[-1].items() if k != 'tokens'}), flush=True)
+    del batcher, future
+    torch.cuda.empty_cache()
+output.write_text(json.dumps(dict(complete=True, rows=rows), indent=2)+'\n')

opt_prime/examples/pd/http_matrix/evidence261/diagnose_local201_koni.py

@@ -0,0 +1,30 @@
+"""Separate GPU batching/chunk effects from a retained exact RNGD KV export."""
+import json
+from pathlib import Path
+from concurrent.futures import Future
+import torch
+from transformers import AutoModelForCausalLM
+from opt_prime.accel.pd import PDRequest
+from opt_prime.accel.pd.batching import ContinuousBatcher
+from opt_prime.accel.pd.pd_batching import PDBatcher,PrefillOut
+r=Path(__file__).resolve().parent
+output=r/'local201-koni-diagnostic.json'
+if output.exists(): raise FileExistsError(output)
+torch.set_num_threads(4)
+e=torch.load(r/'rngd-koni-8b-8192-parts221-layout.pt',weights_only=True,mmap=True)
+d=json.loads((r/'local201-koni-matrix-a.json').read_text())
+assert e['ids']==d['prompts']['8192']
+model=AutoModelForCausalLM.from_pretrained('/tmp/serving-matrix-models/KONI-8B',dtype=torch.bfloat16,attn_implementation='sdpa',local_files_only=True).cuda().eval()
+class Replay:
+ def submit(self,request,send_from=0,slot=None):
+  assert list(request.input_ids)==e['ids']
+  f=Future();f.set_result(PrefillOut(int(e['logits'].argmax()),send_from,e['keys'][:,:,send_from:],e['values'][:,:,send_from:]));return f
+rows=[]
+for mode in ['native256','native8192','retained_rngd_kv']:
+ kw=dict(slots=4,max_context=9216,kv='paged',page_size=64,prefix_cache=False)
+ b=PDBatcher(model,Replay(),**kw) if mode=='retained_rngd_kv' else ContinuousBatcher(model,prefill_chunk=int(mode[6:]),**kw)
+ f=b.submit(PDRequest(mode,tuple(e['ids']),128));b.run_until_idle();tokens=f.result()
+ native=d['references']['8192:128']['tokens'];live=next(x['rows'][0]['tokens'] for x in d['correctness_batches'] if x['input_tokens']==8192 and x['output_tokens']==128)
+ row=dict(mode=mode,tokens=tokens,equals_native=tokens==native,equals_live=tokens==live,first_diff_native=next((i for i,(a,z) in enumerate(zip(tokens,native)) if a!=z),None));rows.append(row);print(json.dumps(row),flush=True)
+ del b;torch.cuda.empty_cache()
+(r/'local201-koni-diagnostic.json').write_text(json.dumps(rows,indent=2)+'\n')