Add head_dim = 64 in B200 Attention.

Summary:
X-link: facebookresearch/FBGEMM#1942

This diff adds support for head_dim = 64 in the B200 Attention module.

Reviewed By: sryap

Differential Revision: D82996471

Author: gchalump

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_bwd.cu

@@ -5,6 +5,7 @@
 template <
     typename Element,
     typename ActiveMask,
+    int HeadDim,
     bool kIsVarlen,
     bool kIsDeterministic,
     class... KernelOptions>
@@ -34,8 +35,8 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> fmha_bwd(
     cute::tuple<VariableLength, VariableLength, int, int, cute::tuple<cute::tuple<int, int>, int>>,
     cute::tuple<int, int, int, int, cute::tuple<cute::tuple<int, int>, int>>
   >;
-
-  using TileShape = Shape<_128, _128, _128>;
+  using D_H = cute::Int<HeadDim>;
+  using TileShape = Shape<_128, _128, D_H>;
 
   using Operation = cutlass::fmha::device::
       Sm100FmhaBwd<ProblemShapeType, Element, ElementAccumulator, TileShape, /*kIsMla=*/false, ActiveMask, kIsDeterministic>;
@@ -114,7 +115,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> fmha_bwd(
     problem_shape = cute::make_tuple(
         Q, K, D, D, make_shape(make_shape(H_R, H_K), B));
   }
-
+  TORCH_CHECK(D == HeadDim);
   TORCH_CHECK(D % 8 == 0); // Alignment
   if constexpr (!kIsVarlen) {
     // TODO: support Q < 8
@@ -314,13 +315,15 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dispatch_fmha_bwd(
     [&](
       auto element,
       auto element_out,
+      auto head_dim,
       auto varlen,
       auto deterministic,
       auto mask,
       auto... kernel_options) {
       return fmha_bwd<
         decltype(element),
         decltype(mask),
+        head_dim,
         varlen,
         deterministic,
         decltype(kernel_options)...>
@@ -340,51 +343,63 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dispatch_fmha_bwd(
         window_size_right);
     };
 
-  auto dispatch_type = [&](auto varlen, auto deterministic, auto mask) {
+  auto dispatch_type = [&](auto varlen, auto deterministic, auto mask, auto head_dim) {
     if (query.dtype() == torch::kFloat16) {
       return dispatch_fmha(
-          cutlass::half_t{}, cutlass::half_t{}, varlen, deterministic, mask);
+          cutlass::half_t{}, cutlass::half_t{}, head_dim, varlen, deterministic, mask);
     }
     else if (query.dtype() == torch::kBFloat16) {
       return dispatch_fmha(
-          cutlass::bfloat16_t{}, cutlass::bfloat16_t{}, varlen, deterministic, mask);
+          cutlass::bfloat16_t{}, cutlass::bfloat16_t{}, head_dim, varlen, deterministic, mask);
     }
     else if (query.dtype() == torch::kFloat8_e4m3fn) {
       return dispatch_fmha(
-          cutlass::float_e4m3_t{}, cutlass::bfloat16_t{}, varlen, deterministic, mask);
+          cutlass::float_e4m3_t{}, cutlass::bfloat16_t{}, head_dim, varlen, deterministic, mask);
     }
     TORCH_CHECK(false, "Unsupported dtype for q: ", query.dtype());
   };
 
+  auto dispatch_head_dim = [&](auto varlen, auto deterministic, auto mask) {
+    if (query.size(query.dim() - 1) == 128) {
+      return dispatch_type(varlen, deterministic, mask, std::integral_constant<int, 128>{});
+    }
+    else if (query.size(query.dim() - 1) == 64) {
+      return dispatch_type(varlen, deterministic, mask, std::integral_constant<int, 64>{});
+    }
+    else {
+      TORCH_CHECK(false, "Unsupported head dim: ", query.size(query.dim() - 1));
+    }
+  };
+
   auto dispatch_mask = [&](auto varlen, auto deterministic) {
     if (causal) {
       if (bottom_right) {
-        return dispatch_type(
+        return dispatch_head_dim(
             varlen, deterministic, CausalForBackwardMask</*kIsQBegin=*/false>{});
       }
       else {
-        return dispatch_type(
+        return dispatch_head_dim(
             varlen, deterministic, CausalForBackwardMask</*kIsQBegin=*/true>{});
       }
     }
     else if (local) {
       if (bottom_right) {
-        return dispatch_type(
+        return dispatch_head_dim(
             varlen, deterministic, LocalMaskForBackward</*kIsQBegin=*/false>{});
       }
       else {
-        return dispatch_type(
+        return dispatch_head_dim(
             varlen, deterministic, LocalMaskForBackward</*kIsQBegin=*/true>{});
       }
     }
     else if (varlen || key.size(1) % 128 != 0) {
       // Use the residual mask for varlen or when K seqlen is not multiple of
       // blockN
-      return dispatch_type(
+      return dispatch_head_dim(
           varlen, deterministic, ResidualMaskForBackward{});
     }
     else {
-      return dispatch_type(
+      return dispatch_head_dim(
           varlen, deterministic, NoMask{});
     }
   };

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_fwd.cu

@@ -320,6 +320,9 @@ std::tuple<at::Tensor, at::Tensor> dispatch_fmha_fwd(
     if (q.size(q.dim() - 1) == 128) {
       return dispatch_type(varlen, mask, std::integral_constant<int, 128>{});
     }
+    else if (q.size(q.dim() - 1) == 64) {
+      return dispatch_type(varlen, mask, std::integral_constant<int, 64>{});
+    }
     else {
       TORCH_CHECK(false, "Unsupported head dim: ", q.size(q.dim() - 1));
     }

fbgemm_gpu/experimental/gen_ai/test/attention/blackwell_fmha_test.py

@@ -440,17 +440,26 @@ def _execute_cutlass_blackwell_attn_varlen(
                 seqlen_k,
                 batch_size,
                 is_mqa,
+                window_size,
+                head_dim,
+                sm_scale,
             )
             for seqlen_k in [64, 128, 256, 1024]
             for batch_size in [1, 2]
             for is_mqa in [True]
+            for window_size in [(-1, -1), (0, 0), (0, 128), (128, 0), (1024, 0)]
+            for head_dim in [128]
+            for sm_scale in [None, 1.0 / head_dim]
         ]
     )
     def test_decode(
         self,
         seqlen_k: int,
         batch_size: int,
         is_mqa: bool,
+        window_size: tuple[int, int],
+        head_dim: int,
+        sm_scale: Optional[float],
         q_heads: int = 8,
         dtype: torch.dtype = torch.float8_e4m3fn,
     ) -> None:
@@ -465,7 +474,7 @@ def test_decode(
             seqlen_k,
             q_heads,
             kv_heads=1 if is_mqa else q_heads,
-            head_dim=128,
+            head_dim=head_dim,
             dtype=dtype,
             causal=causal,
             # Decode kernel does not support sliding window attention yet
@@ -486,14 +495,16 @@ def test_decode(
                 q_heads,
                 causal,
                 window_size,
+                head_dim,
                 sm_scale,
             )
             for kv_padding in [128, 256, 512, 1024]
             for batch_size in [2, 8]
             for q_heads in [8, 16]
             for causal in [True, False]
             for window_size in [(-1, -1), (0, 0), (0, 128), (128, 0), (1024, 0)]
-            for sm_scale in [None, 1.0 / 128]
+            for head_dim in [128]
+            for sm_scale in [None, 1.0 / head_dim]
         ]
     )
     def test_jagged_vs_padded_kv(
@@ -503,6 +514,7 @@ def test_jagged_vs_padded_kv(
         q_heads: int,
         causal: bool,
         window_size: tuple[int, int],
+        head_dim: int,
         sm_scale: Optional[float],
     ) -> None:
         """
@@ -517,7 +529,7 @@ def test_jagged_vs_padded_kv(
         seqlen_q = kv_padding  # Maximum sequence length (padded size)
         device = torch.accelerator.current_accelerator()
         kv_heads = 1
-        head_dim = 128
+        head_dim = head_dim
         dtype = torch.bfloat16
 
         torch.manual_seed(SEED)
@@ -663,6 +675,7 @@ def test_jagged_vs_padded_kv(
                 is_varlen,
                 kv_heads,
                 window_size,
+                head_dim,
                 sm_scale,
             )
             for seqlen_q, offset_q in [
@@ -682,7 +695,8 @@ def test_jagged_vs_padded_kv(
             for is_varlen in [False, True]
             for kv_heads in [1, 2, 3, 4]
             for window_size in [(-1, -1), (0, 0), (0, 128), (128, 0), (1024, 0)]
-            for sm_scale in [None, 1.0 / 128]
+            for head_dim in [64, 128]
+            for sm_scale in [None, 1.0 / head_dim]
         ]
     )
     def test_forward(
@@ -695,6 +709,7 @@ def test_forward(
         is_varlen: bool,
         kv_heads: int,
         window_size: tuple[int, int],
+        head_dim: int,
         sm_scale: Optional[float],
         dtype: torch.dtype = torch.bfloat16,
     ) -> None:
@@ -713,7 +728,7 @@ def test_forward(
             seqlen_k,
             q_heads=kv_heads * q_heads_per_kv_head,
             kv_heads=kv_heads,
-            head_dim=128,
+            head_dim=head_dim,
             dtype=dtype,
             causal=causal,
             window_size=window_size,
@@ -736,7 +751,8 @@ def test_forward(
             [(-1, -1), (128, 0), (256, 0), (128, 128), (512, 0)]
         ),
         deterministic=st.booleans(),
-        sm_scale=st.sampled_from([None, 1.0 / 128]),
+        head_dim=st.sampled_from([64, 128]),
+        is_sm_scale=st.booleans(),
     )
     @settings(**common_settings)
     def test_backward(
@@ -750,8 +766,10 @@ def test_backward(
         is_gqa: bool,
         window_size: tuple[int, int],
         deterministic: bool,
-        sm_scale: Optional[float],
+        head_dim: int,
+        is_sm_scale: bool,
     ) -> None:
+        sm_scale = 1.0 / head_dim if is_sm_scale else None
         test_func = (
             self._execute_cutlass_blackwell_attn_varlen
             if is_varlen
@@ -764,7 +782,7 @@ def test_backward(
             seqlen,
             q_heads=kv_heads * q_heads_per_kv_head,
             kv_heads=kv_heads,
-            head_dim=128,
+            head_dim=head_dim,
             dtype=dtype,
             causal=causal,
             window_size=window_size,