revert disable flash operator on rocm

tenpercent · Aug 16, 2024 · d6b6456 · d6b6456
1 parent 7d21800
commit d6b6456
Showing 1 changed file with 2 additions and 8 deletions.
diff --git a/xformers/ops/fmha/flash.py b/xformers/ops/fmha/flash.py
@@ -607,10 +607,7 @@ class FwOp(AttentionFwOpBase):
         implementation.
     """
 
-    if torch.version.hip:
-        OPERATOR = None
-    else:
-        OPERATOR = get_operator("xformers_flash", "flash_fwd")
+    OPERATOR = get_operator("xformers_flash", "flash_fwd")
     SUPPORTED_DEVICES: Set[str] = {"cuda"}
     CUDA_MINIMUM_COMPUTE_CAPABILITY = (8, 0)
     SUPPORTED_DTYPES: Set[torch.dtype] = {torch.half, torch.bfloat16}
@@ -812,10 +809,7 @@ def operator_flop(
 class BwOp(AttentionBwOpBase):
     __doc__ = FwOp.__doc__
 
-    if torch.version.hip:
-        OPERATOR = None
-    else:
-        OPERATOR = get_operator("xformers_flash", "flash_bwd")
+    OPERATOR = get_operator("xformers_flash", "flash_bwd")
     SUPPORTED_DEVICES = FwOp.SUPPORTED_DEVICES
     CUDA_MINIMUM_COMPUTE_CAPABILITY = FwOp.CUDA_MINIMUM_COMPUTE_CAPABILITY
     SUPPORTED_DTYPES = FwOp.SUPPORTED_DTYPES