File size: 3,558 Bytes
105f9ef d144fee 105f9ef d144fee 105f9ef d144fee 105f9ef d144fee 105f9ef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """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()
|