migrate "jagged_flash_attention"

Summary: As titled.

Differential Revision: D66932340

fbgemm_gpu/fbgemm_gpu/sll/cpu_sll.py

@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
-from typing import Any, Tuple
+from typing import Any, Callable, Tuple
 
 import torch