Add python API for flash-attn#558

YangKai0616 wants to merge 2 commits intointel:mainintel/sycl-tla:mainfrom

YangKai0616:mainYangKai0616/sycl-tla-pythonAPI-flash-attn:mainCopy head branch name to clipboard