[TLX] [FA] Enable H-DIM=64 for Backwards (#711)

Summary:
Enables H-DIM=64 for Backwards which already works.

Pull Request resolved: #711

Reviewed By: htyu

Differential Revision: D88212441

Pulled By: njriasan

fbshipit-source-id: f90d0b4fc652002a6ab95f39abd423e5eb9668c6

Author: njriasan

third_party/tlx/tutorials/blackwell-fa-ws-pipelined-persistent_test.py

@@ -1750,11 +1750,7 @@ def test_op(Z, H, N_CTX, HEAD_DIM, mode, provider, causal, dtype=torch.float16):
     sm_scale = 0.5
     # reference implementation
     ref_dtype = dtype
-    if mode == "bwd" and HEAD_DIM == 64:
-        pytest.skip("Only test bwd with 128")
-    elif mode == "fwd" and not causal and HEAD_DIM == 128:
-        pytest.skip("Only test fwd with causal")
-    elif mode == "bwd" and causal:
+    if mode == "bwd" and causal:
         pytest.skip("Causal not supported for bwd yet")
     if mode == "fwd" and "fp8" in provider:
         ref_dtype = torch.float32