| """Owner-selected bf16 segmented attention and shared-prefix inference.""" |
| import os |
| import torch |
| import saif_v2_core as core |
| from transformers.models.qwen3_5 import modeling_qwen3_5 as modeling |
|
|
| MASK=core.attention_mask |
| ATTENTION=modeling.ALL_ATTENTION_FUNCTIONS['sdpa'] |
| LAYOUT=None |
|
|
|
|
| def mask(layout, attention, device, dtype, boolean=False): |
| global LAYOUT |
| assert attention=='bidir' |
| LAYOUT=layout |
| return {'full_attention':None,'linear_attention':None} |
|
|
|
|
| def segmented(module,query,key,value,attention_mask,dropout=0.,scaling=None,**kwargs): |
| assert dropout==0 and LAYOUT is not None |
| prefix=LAYOUT.prefix |
| assert query.shape[-2]==len(LAYOUT.ids) |
| if module.num_key_value_groups>1: |
| key=key.repeat_interleave(module.num_key_value_groups,dim=1) |
| value=value.repeat_interleave(module.num_key_value_groups,dim=1) |
| def attend(q,k,v): |
| return torch.nn.functional.scaled_dot_product_attention(q,k,v,dropout_p=0.,is_causal=False,scale=scaling) |
| output=[attend(query[:,:,:prefix],key[:,:,:prefix],value[:,:,:prefix])] |
| for start,end in LAYOUT.blocks: |
| output.append(attend(query[:,:,start:end],torch.cat((key[:,:,:prefix],key[:,:,start:end]),dim=2),torch.cat((value[:,:,:prefix],value[:,:,start:end]),dim=2))) |
| return torch.cat(output,dim=2).transpose(1,2).contiguous(),None |
|
|
|
|
| def install(): |
| core.attention_mask=mask |
| modeling.ALL_ATTENTION_FUNCTIONS['sdpa']=segmented |
| os.environ['SAIF_EXPERIMENTAL_SHARED_DELTA']='1' |
|
|
|
|
| def uninstall(): |
| core.attention_mask=MASK |
| modeling.ALL_ATTENTION_FUNCTIONS['sdpa']=ATTENTION |
| os.environ.pop('SAIF_EXPERIMENTAL_SHARED_DELTA',None) |
|
|
