/jKUdZddlmZddlZddlZddlmZerddlmZddl m Z ddl m Z ddl mZmZddlZdd lmZd d lmZd gZdad ed<dad ed<e GddZe dadZ dbdcdZdddZdedZdfd+Zdgd2Zdhd9Z ed:Z!did=Z"djd?Z# dkdldHZ$ dmdndNZ% doddOdOddddPdpdRZ&ddOdOddddPdqdSZ'ddddTdrdWZ( dsddYdtdZZ) duddYdvd[Z*ddYdwd^Z+ej,d_e`dS)xz PROTOTYPE! Flash Attention 3 implementation. For fp8: only supports forward pass right now. For fp16/bf16: supports forward and backward pass. ) annotationsN) TYPE_CHECKING)Callable) dataclass)cache) TypeVarTupleUnpack)Library) _registryregister_flash_attention_fa3zCallable | None _FA3_CUDA_FWD _FA3_CUDA_BWDc"eZdZUded<ddZdS) _FA3HandlezLibrary | NonelibraryreturnNonecRd|_tjddS)NF)rtorch_C_set_sdp_use_fa3)selfs \/home/longshao/multi-rider-rag/.venv/lib/python3.11/site-packages/torch/nn/attention/_fa3.pyremovez_FA3Handle.remove*s%  !!%(((((N)rr)__name__ __module__ __qualname____annotations__rrrrr&s6))))))rrdevice torch.devicerintcJtj|\}}|SN)rcudaget_device_capability)r"major_s r_get_device_majorr+0s z//77HE1 Lrflash_attn_interface module_pathstrct|tjdt t S)z Register FA3 flash attention kernels with the PyTorch dispatcher. Args: module_path: Python module path to the FA3 implementation. T)_fa3_import_modulerrrr_fa3_register_kernelsr-s rr r 6s?{### Hd### +-- . ..rrctj|ttjdst d|dttjjdst d|dttjjdst d|dtjjjatjjj a dS)N flash_attn_3zModule 'z' does not expose FA3 kernelsfwdz%' does not expose FA3 forward kernelsbwdz&' does not expose FA3 backward kernels) importlib import_modulehasattrrops RuntimeErrorr4r5rr6rr2s rr0r0Gs K((( 59n - -RPkPPPQQQ 59)5 1 1  I{ I I I    59)5 1 1  J{ J J J   I*.MI*.MMMrr cxtddd}|dtd|dtd|dtd|dt d|dt d|d td|S) NatenIMPLCUDAz"_flash_attention_forward.quantizedz-_scaled_dot_product_flash_attention.quantized_flash_attention_forward#_scaled_dot_product_flash_attention_flash_attention_backward,_scaled_dot_product_flash_attention_backward)r impl!_fa3_flash_attention_forward_impl4_fa3_scaled_dot_product_flash_attention_forward_impl)_fa3_flash_attention_forward_impl_default<_fa3_scaled_dot_product_flash_attention_forward_impl_default"_fa3_flash_attention_backward_impl5_fa3_scaled_dot_product_flash_attention_backward_impl)libs rr1r1Xs &&& ) )CHH,.OQWHH7< HH"$MvHH-D HH (*LfUUUHH6= Jrquery torch.Tensortensorstuple[torch.Tensor, ...] dropout_pfloat cum_seq_qtorch.Tensor | None q_descale k_descale v_descale str | Nonec|dkrdStd|DsdStd|DdkrdS|jtjkr |||t jdt||d krd S||d krd Stj sd St|j dkrdSdS)Nzdropout_p must be 0c3$K|] }|jV dSr&)is_cuda.0ts r z,_fa3_common_support_error..s$**Qqy******rzinputs must be CUDA tensorsch|] }|j Sr!)r"r\s r z,_fa3_common_support_error..s & & &AH & & &rr inputs must share devicezWhen using SDPA with fp8, descale tensor should always be used for accurate dequantization. Please use _scaled_dot_product_attention_quantized and provide the descale tensors.zdense query must be 4Dzragged query must be 3DzCUDA not available z#FA3 requires compute capability 9.0) alllendtyper float8_e4m3fnwarningswarn UserWarningdimr' is_availabler+r")rLrNrPrRrTrUrVs r_fa3_common_support_errorrotsC$$ **'*** * *-,, & &g & & &''1,,)) {e)))Y.)2C  +     UYY[[A--''!1!1(( : " " $ $$##&&!++44 4rkeyvaluereturn_debug_maskbool alibi_slopes seqused_kc t |rdS|dS| |jtjkrdS|jsdStjtjtjf t fd|||hDsd Std|||hDdkrd St||||f|||| | } | | d krd S| SdS) Nzreturn_debug_mask must be Falsezalibi_slopes not supportedzseqused_k must be int32zseqused_k must be CUDAc3*K|] }|jvVdSr&rhr]r^supported_dtypess rr_z-_fa3_forward_support_error..s+HHqqw**HHHHHHrinputs must be one of ch|] }|j Sr!rxr\s rraz-_fa3_forward_support_error..s 1 1 1AG 1 1 1rr #all inputs must have the same dtyperbz(query, key, value must be on same device) rhrint32r[rifloat16bfloat16rfrgro) rLrprqrPrrrtrurRrTrUrVerrorrzs @r_fa3_forward_support_errorrs100++ ?ek ) ),,  ,+++U]ENK HHHHUC4GHHH H H;:(8::: 1 1eS%0 1 1 122a7744 % U  E  . . .== 4rgrad_outout logsumexpwindow_size_left int | Nonewindow_size_rightc f |jtjkr dS|jtjkrdStjtjf t fd|||||hDsd Std|||||hDdkrdSt|||||||f||ddd} | | SdS)NzHFA3 backward does not support fp8 - use inference only (torch.no_grad())zlogsumexp dtype must be float32c3*K|] }|jvVdSr&rxrys rr_z._fa3_backward_support_error..s+WWqqw**WWWWWWrr{ch|] }|j Sr!rxr\s rraz._fa3_backward_support_error..s @ @ @AG @ @ @rr r}) rhrrifloat32rrrfrgro) rrLrprqrrrPrRrrrrzs @r_fa3_backward_support_errorrs {e))) V  %-''00 u~6 WWWWXuc5RU4VWWW W W;:(8::: @ @hsE3? @ @ @AAQFF44 % 5#uc95     E  4rTs Unpack[Ts]tuple[Unpack[Ts]]c4td|DS)Nc3BK|]}|ddVdS)r N) transposer\s rr_z#_transpose_dense..s044qQ""444444r)tuple)rNs r_transpose_densers 44G444 4 44rxcd|-|ddkr|n|S)z2Ensure tensor is contiguous in the last dimension.Nr )stride contiguous)rs r_maybe_contiguousrs,]qxx||q/@/@1<<>>>aGrcu_seq_qcu_seq_kmax_qmax_kscale float | None is_causal!tuple[torch.Tensor, torch.Tensor]cttdt|}t|}|jtjkrF|ddkr-|ddkr|nt|}t|}t|}t| } tg|||ddd| ||dd| ||dddddd| ||||| | nd| | ndddddt jrdnddtj pdR\} }}}| |fS) zF Run the FA3 forward pass by calling the C++ kernel directly. NFA3 not registeredrr rrYT) rr;rrhrrirr$are_deterministic_algorithms_enabledr_get_sm_carveout_experimental)rLrprqrrrrrrrrrurrTrUrVqkv cu_seqlens_q cu_seqlens_k softmax_lse out_accumsoftmax_lse_accums r_fa3_run_forwardrs*/000%  A#A ;%- - - LL   ! ! LL   ! ! u % % %X..L$X..L!),,I5B#6 #6 #6 #6 #6 #6 #6 #6 #6 #6 #6 #6 #6 #6 #6 #6 !#6" ##6$ %#6& '#6( )#6* +#6, -#6. /#60 1#62 3#64-8b5#66/:7#68 9#6: ;#6< =#6> ?#6@ 7 9 9@qA#6B C#6D ..005AE#6#6#62Ci!2H  &&(( ((rF max_seqlen_q max_seqlen_k deterministic/tuple[torch.Tensor, torch.Tensor, torch.Tensor]cttdt|}|ddkr|n|}|ddkr|n|}|ddkr|n|}t|}t|}t j|}t j|}t j|}t|||||||||||dd|| | | | | d|t jpd|||fS)Nrrr rYr) rr;rrrr empty_likerr)rrLrprqrrrrrrrrrrrdoutrrrolsedqdkdvs r_fa3_run_backwardrCsi"/000 X & &D#ll2..!33AJJrNNa//SA#ll2..!33A#A I & &C  !  B  !  B  !  B              ..005A-0 r2:rrrrrrurtr cum_seq_kc t||||| |||| | | }|td|t|||||||| |||||| | | \}}tjdtj|j}tjdtj|j}tjd|j|j}|||||fS)Nz)FA3 flash_attention forward unsupported: )r)rhr"r!r) rr;rrzerosuint64r"emptyrh)rLrprqrRrrrrPrrrrTrUrVrrrrurtrrr rng_state philox_offset debug_masks rrErE~s, '     E NuNNOOO       !HC$ D U\JJJIK%,u|LLLMQek%,GGGJ Y z 99rc Ft|||||||||| ddd| | | | ||S)Nr)rE)rLrprqrRrrrrPrrrrrrrurtrs rrGrGsT& -        )+! '   r)rrrrunusedct||||||| ||| }|td|tj}t |||||||||| || ||nd||nd|\}}}|||fS)z0FA3 implementation of _flash_attention_backward.Nz*FA3 flash_attention backward unsupported: r)rr;rrr)rrLrprqrrrRrrrrPrrrrrrrrrrrs rrIrIs* (      E OOOPPP>@@M"       ,8b.:JBB" r2:rrYrc 4t|||||ddd||| } | td| t|||\} } } |jtjkr tjn|j}t j||}|dd}| d}| d}t| | | dd|||||| ||||\}}}}}| d}| d}||dd|||||f S)NzFA3 SDPA forward unsupported: rxr r)rrrTrUrV) rr;rrhrrirrrsizerE)rLrprqrTrUrVrPrrrrrrrr out_dtypeout_bhsdout_bshd max_q_flash max_k_flashr*rrrrrrs rrFrF sa '        E CECCDDDuc511GAq! #(+1D"D"D%+IY777H!!!Q''H&&))K&&))K3T      4440AsI}j" JJqMME HHQKKE       rc 4t|||ddd|||| S)Nr)rF)rLrprqrPrrrrs rrHrHgs: @          r philox_seedrct||||||| ddd }|td|t|||||\}}}}}t||||||dd|| | | | | |\}}}t|||\}}}|||fS)zCFA3 implementation of _scaled_dot_product_flash_attention_backward.NzFA3 SDPA backward unsupported: r)rr;rrI)rrLrprqrrrRrrrrPrrrrr grad_out_tq_tk_tv_tout_trrrdq_outdk_outdv_outs rrJrJs& (%eS)YdD  E DUDDEEE(8%eS(($JS#u4        JBB&.b"b99FFF 66 !!rFA3) register_fn)r"r#rr$)r,)r-r.rr)r-r.rr)rr )rLrMrNrOrPrQrRrSrTrSrUrSrVrSrrW)rLrMrprMrqrMrPrQrrrsrtrSrurSrRrSrTrSrUrSrVrSrrW)rrMrLrMrprMrqrMrrMrrMrPrQrRrSrrrrrrW)rNrrr)rrSrrS)NNNN)"rLrMrprMrqrMrrSrrSrr$rr$rrrrsrrrrrurSrrSrTrSrUrSrVrSrr)F) rrMrLrMrprMrqrMrrMrrMrrSrrSrrrrrrrrsrr$rr$rrsrr)NNN)&rLrMrprMrqrMrRrSrrSrr$rr$rPrQrrsrrrsrTrSrUrSrVrSrrrr$rr$rurSrtrSrrS) rLrMrprMrqrMrRrSrrSrr$rr$rPrQrrsrrrsrrrr$rr$rurSrtrSrrS)"rrMrLrMrprMrqrMrrMrrMrRrSrrSrr$rr$rPrQrrsrrMrrMrrrrrr)NNNrYFF)rLrMrprMrqrMrTrSrUrSrVrSrPrQrrsrrrsrr)rYFF)rLrMrprMrqrMrPrQrrsrrrsrr)rrMrLrMrprMrqrMrrMrrMrRrSrrSrr$rr$rPrQrrsrrMrrMrr)-__doc__ __future__rr7rjtypingrcollections.abcr dataclassesr functoolsrtyping_extensionsrr r torch.libraryr r __all__rr rrr+r r0r1rorrrrrrrrErGrIrFrHrJregister_flash_attention_implr!rrrs#""""" )((((((!!!!!!22222222 !!!!!!#  "& %%%%!% %%%% ))))))) ) ./////"////"8""""J((((V####L\$5555HHHH$ $%)%)%)!J)J)J)J)J)x 88888L&*%)%)::%)(,#)::::::::::::R%)(,##''''''t#'$(%888888~&*%)%)#DDDDDDDV# P!2"2"2"2"2"2"j( ';WXXXXXXr