"""Flash attention forward task."""