"""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:])