Fix performance regression in cutlass backward pass since 0.0.25 (#1177)

Fixes #1176
facebookresearch · Dec 17, 2024 · b17d794 · b17d794
1 parent 3d947a6
commit b17d794
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/xformers/ops/fmha/cutlass.py b/xformers/ops/fmha/cutlass.py
@@ -436,7 +436,7 @@ def apply(cls, ctx: Context, inp: Inputs, grad: torch.Tensor) -> Gradients:
             philox_offset=rng_offset,
             custom_mask_type=_custom_mask_type(inp.attn_bias),
             scale=inp.scale,
-            num_splits_key=-1,  # Let C++ determine it
+            num_splits_key=None,  # Let C++ determine it
             window_size=(
                 inp.attn_bias._window_size
                 if isinstance(