"""Versioned CUDA repair: bounded FP32 attention, FP64 recurrent accumulation.
No changes to the deployed v1 engine or existing evidence.

v1.1 (Solomon v1.1): the repair is the fp32 precision only. Under bf16 / int8 the engine is exactly the
'bf16-fp32-recurrence' mode the trainer runs in (bf16 SDPA attention, fp32 recurrence), and the identity says
so; the q-block attention here computes in the query dtype, so installing it under bf16 would have been
bf16 attention labelled 'fp32-attention'.
"""
import hashlib,inspect,json,textwrap
from pathlib import Path
from solomon.engine_cuda import CudaEngine as V1

def install_numerics():
    import torch
    from transformers.models.qwen3_5 import modeling_qwen3_5 as m
    # Use the pinned readable reference implementation rather than FLA's FP32
    # chunk solve. Output retains input dtype; only recurrence arithmetic is FP64.
    for name in ('torch_chunk_gated_delta_rule','torch_recurrent_gated_delta_rule'):
        original=inspect.unwrap(getattr(m,name))   # promoted wrappers carry __wrapped__, so this is the reference
        if getattr(original,'_solomon_promoted',False):raise ValueError('recurrence reference not reachable')
        source=textwrap.dedent(inspect.getsource(original))
        source=source[source.index('def '):].replace('torch.float32','torch.float64')
        namespace=dict(vars(m));exec(compile(source,'<qwen-fp64-reference>','exec'),namespace)
        setattr(m,name,namespace[name])
    def bounded(module,query,key,value,attention_mask,dropout=0.,scaling=None,**kwargs):
        groups=query.shape[1]//key.shape[1]
        key=m.repeat_kv(key,groups);value=m.repeat_kv(value,groups)
        length=query.shape[-2];total=key.shape[-2];outputs=[]
        for start in range(0,length,128):
            end=min(length,start+128)
            scores=query[...,start:end,:]@key.transpose(-2,-1)
            scores=scores*(scaling if scaling is not None else query.shape[-1]**-.5)
            if attention_mask is not None:
                mask=attention_mask[...,start:end,:total] if attention_mask.shape[-2]!=1 else attention_mask[...,:total]
                if mask.dtype==torch.bool:scores=scores.masked_fill(~mask,float('-inf'))
                else:scores=scores+mask
            elif getattr(module,'is_causal',False):
                qpos=torch.arange(start,end,device=query.device)+total-length
                kpos=torch.arange(total,device=query.device)
                scores=scores.masked_fill(kpos[None,:]>qpos[:,None],float('-inf'))
            probs=scores.softmax(-1)
            outputs.append(probs@value)
        return torch.cat(outputs,dim=-2).transpose(1,2).contiguous(),None
    # Replace the decoder's existing SDPA interface. Vision retains its installed
    # module interface only insofar as it also dispatches here; causal flag false
    # gives ordinary dense vision attention without a causal mask.
    m.ALL_ATTENTION_FUNCTIONS.register('sdpa',bounded)

class CudaEngine(V1):
    def __init__(self,*args,**kwargs):
        super().__init__(*args,**kwargs)
        if self.mode=='fp32':install_numerics()
        self.identity['base_engine_sha256']=self.identity['code_sha256']
        self.identity['code_sha256']=hashlib.sha256(Path(__file__).read_bytes()).hexdigest()
        if self.mode=='fp32':self.identity['numerics']='fp32-attention-qblock128-fp64-recurrence-reference-v2'
        self.identity.pop('fingerprint',None)
        self.identity['fingerprint']=hashlib.sha256(json.dumps(self.identity,sort_keys=True).encode()).hexdigest()