| """Segment-preserving PyTorch fallback for CPU FlashAttention calls.""" |
| import torch |
| from torch.nn import functional as F |
| def flash_attn_varlen_func(q,k,v,cu_seqlens_q,cu_seqlens_k,max_seqlen_q,max_seqlen_k,dropout_p=0.,softmax_scale=None,causal=False,**kwargs): |
| output=torch.empty_like(q) |
| for i in range(len(cu_seqlens_q)-1): |
| a,b=int(cu_seqlens_q[i]),int(cu_seqlens_q[i+1]);c,d=int(cu_seqlens_k[i]),int(cu_seqlens_k[i+1]) |
| if a==b:continue |
| qq=q[a:b].transpose(0,1).unsqueeze(0);kk=k[c:d].transpose(0,1).unsqueeze(0);vv=v[c:d].transpose(0,1).unsqueeze(0) |
| y=F.scaled_dot_product_attention(qq,kk,vv,dropout_p=dropout_p,is_causal=causal,scale=softmax_scale,enable_gqa=qq.shape[1]!=kk.shape[1]);output[a:b]=y.squeeze(0).transpose(0,1) |
| return output |
| def index_first_axis(x,indices):return x[indices] |
| def unpad_input(x,attention_mask): |
| lengths=attention_mask.sum(-1,dtype=torch.int32);indices=torch.nonzero(attention_mask.flatten(),as_tuple=False).flatten();cu=F.pad(lengths.cumsum(0,dtype=torch.int32),(1,0));return x.flatten(0,1)[indices],indices,cu,int(lengths.max()) |
| def pad_input(x,indices,batch,seqlen): |
| y=x.new_zeros((batch*seqlen,*x.shape[1:]));y[indices]=x;return y.reshape(batch,seqlen,*x.shape[1:]) |
|
|
