"""Architecture-derived Q/K gain consolidation on non-RoPE dimensions only. Fixed layer3 interventions; no fitting, layer selection or learned alignment. Run only after baseline/candidate pilot finishes, within the same cumulative budget. """ import argparse import csv import json from pathlib import Path import resource import time import numpy as np from .model import StreamingLM,norm,rope from .run import probability_metrics,write_json,save_csv LAYERS=[3,7,11,15,19,23] def normalized(x): return x/np.sqrt(np.mean(x*x,axis=-1,keepdims=True)+np.float32(1e-6)) def consolidate(q,k,a,b,r,move,joint=None): q,k=normalized(q),normalized(k) c=a[r:]*b[r:] if joint is None else joint if move=='q_to_k': q=np.concatenate([q[...,:r]*a[:r],q[...,r:]],axis=-1) k=np.concatenate([k[...,:r]*b[:r],k[...,r:]*c],axis=-1) else: q=np.concatenate([q[...,:r]*a[:r],q[...,r:]*c],axis=-1) k=np.concatenate([k[...,:r]*b[:r],k[...,r:]],axis=-1) return q,k class FactoredLM(StreamingLM): def configure(self,layer,move): self.factor_layer=layer;self.move=move p=f'layers.{layer}.self_attn.' self.a=1+self.small(p+'q_norm.weight');self.b=1+self.small(p+'k_norm.weight') self.r=int(self.config['head_dim']*self.config['rope_parameters']['partial_rotary_factor']) self.joint=self.a[self.r:]*self.b[self.r:] def normalization(self,x,name): p=f'layers.{self.factor_layer}.self_attn.' if name not in (p+'q_norm',p+'k_norm'):return super().normalization(x,name) # This path does not multiply the unit-gain tail, unlike a full-vector mask. unit=x/np.sqrt(np.mean(x*x,axis=-1,keepdims=True)+self.eps) isq=name==p+'q_norm';head=self.a if isq else self.b skip=(isq and self.move=='q_to_k') or (not isq and self.move=='k_to_q') tail=unit[...,self.r:] if skip else unit[...,self.r:]*self.joint return np.concatenate([unit[...,:self.r]*head[:self.r],tail],axis=-1) def main(): p=argparse.ArgumentParser();p.add_argument('--root',required=True,type=Path);a=p.parse_args();root=a.root # Protocol is fixed before inspecting functional measurements. protocol={'fixed_intervention_layer':3,'directions':['q_to_k','k_to_q'],'diagnostic_layers':LAYERS,'no_selection':True,'no_training':True,'precision':'FP32 product; no BF16 requantization','quality_gate':{'mean_nll_delta':.01,'sequence_p95_nll_delta':.05}} protocol_path=root/'gain-factorization-protocol.json' if protocol_path.exists(): old=json.loads(protocol_path.read_text());assert all(old[k]==v for k,v in protocol.items()), 'Protocol differs from preregistration' else:write_json(protocol_path,{**protocol,'registered_unix':time.time()}) ledger_path=root/'compute-budget-ledger.json';ledger=json.loads(ledger_path.read_text()) remaining=int(7200-ledger['spent_seconds']-30) if remaining<60:raise TimeoutError('No remaining cumulative budget for deterministic identity verification') start=time.monotonic() report={'state':'running','records':[],'interventions':[]} try: cfg=json.loads((root/'model_config.json').read_text());c=cfg['text_config'];r=int(c['head_dim']*c['rope_parameters']['partial_rotary_factor']) assert c['rms_norm_eps']==1e-6 assert c['head_dim']==256 and r==64 and c['num_attention_heads']==8 and c['num_key_value_heads']==2 lm=FactoredLM(root/'weights',cfg,seconds=remaining) data=np.load(root/'baseline-activations.npz');samples=json.loads((root/'samples-tokenized.json').read_text())['samples'] rp=c['rope_parameters'];repeats=c['num_attention_heads']//c['num_key_value_heads'] records=[] for layer in LAYERS: lm.guard() pfx=f'layers.{layer}.self_attn.' aw=lm.small(pfx+'q_norm.weight');bw=lm.small(pfx+'k_norm.weight');ag,bg=1+aw,1+bw q,k=data[f'qraw_{layer}'],data[f'kraw_{layer}'] qr,kr=rope(norm(q,aw),rp['rope_theta'],rp['partial_rotary_factor']),rope(norm(k,bw),rp['rope_theta'],rp['partial_rotary_factor']) ref=np.einsum('bthd,bshd->bhts',qr,np.repeat(kr,repeats,axis=2))/np.float32(np.sqrt(c['head_dim'])) mask=np.triu(np.ones(ref.shape[-2:],dtype=bool),1) def probs(s): s=np.where(mask,-np.inf,s);e=np.exp(s-s.max(-1,keepdims=True));return e/e.sum(-1,keepdims=True) pref=probs(ref) for move in ('q_to_k','k_to_q'): qq,kk=consolidate(q,k,ag,bg,r,move) qq,kk=rope(qq,rp['rope_theta'],rp['partial_rotary_factor']),rope(kk,rp['rope_theta'],rp['partial_rotary_factor']) changed=np.einsum('bthd,bshd->bhts',qq,np.repeat(kk,repeats,axis=2))/np.float32(np.sqrt(c['head_dim'])) pp=probs(changed) for bi,s in enumerate(samples): diff=changed[bi].astype(np.float64)-ref[bi] records.append({'layer':layer,'direction':move,'sample_id':s['id'],'split':s['split'], 'score_max_abs_error':float(np.abs(diff).max()),'score_relative_l2':float(np.linalg.norm(diff)/max(np.linalg.norm(ref[bi].astype(np.float64)),1e-12)), 'probability_max_abs_error':float(np.abs(pp[bi]-pref[bi]).max()),**probability_metrics(pref[bi:bi+1],pp[bi:bi+1])}) save_csv(root/'gain-factorization-score-metrics.csv',records) held=np.array([i for i,s in enumerate(samples) if s['split']=='heldout']);ids=np.array([s['token_ids'] for s in samples],dtype=np.int64) base=np.load(root/'baseline-heldout-token-nll.npy').mean(1) for move in ('q_to_k','k_to_q'): lm.configure(3,move) h,_=lm.forward(ids[held]) nll=lm.nll(h,ids[held]);np.save(root/f'gain-factorization-{move}-token-nll.npy',nll) delta=nll.mean(1)-base report['interventions'].append({'layer':3,'direction':move,'sequence_nll_deltas':delta.tolist(),'mean_nll_delta':float(delta.mean()),'sequence_p95_delta':float(np.quantile(delta,.95)),'passed':bool(delta.mean()<=.01 and np.quantile(delta,.95)<=.05)}) write_json(root/'gain-factorization-summary.json',report) # Local normalization timing uses precomputed gains for both implementations. q,k=data['qraw_3'][:1],data['kraw_3'][:1];pfx='layers.3.self_attn.' ag,bg=1+lm.small(pfx+'q_norm.weight'),1+lm.small(pfx+'k_norm.weight') def baseline():return normalized(q)*ag,normalized(k)*bg joint=ag[r:]*bg[r:] def factored():return consolidate(q,k,ag,bg,r,'q_to_k',joint=joint) for _ in range(3):baseline();factored() timing=[] for repeat in range(10): for label,fn in (('baseline',baseline),('factored',factored)) if repeat%2==0 else (('factored',factored),('baseline',baseline)): lm.guard();t=time.perf_counter() for _ in range(20):fn() timing.append({'repeat':repeat,'implementation':label,'seconds_per_call':(time.perf_counter()-t)/20}) save_csv(root/'gain-factorization-kernel-timing.csv',timing) report.update(state='completed',maximum_score_abs_error=max(x['score_max_abs_error'] for x in records),maximum_probability_abs_error=max(x['probability_max_abs_error'] for x in records), bidirectional_heldout_pass=all(x['passed'] for x in report['interventions']), elementwise_gain_multiplies_saved_per_token_layer={'q_to_k':1536,'k_to_q':384}, original_bf16_payload_bytes=1024,factored_lossless_mixed_payload_bytes=1024, fp32_cached_gain_payload_before=2048,fp32_cached_gain_payload_after=1280, claim_limit='Exact real-arithmetic identity only on non-RoPE coordinates; empirical FP32/NLL checks on tiny text pilot. Not evidence that two existing full layers interchange, not model-wide speedup.') except Exception as exc: report.update(state='stopped',error_type=type(exc).__name__,error=str(exc));raise finally: elapsed=time.monotonic()-start report.update(elapsed_seconds=elapsed,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) ledger=json.loads(ledger_path.read_text());ledger['spent_seconds']+=elapsed;ledger['calls'].append({'module':'functional.gain_factorization','seconds':elapsed});write_json(ledger_path,ledger) write_json(root/'gain-factorization-summary.json',report) if __name__=='__main__':main()