Enable deterministic mode in Cutlass attn (#4840)

Summary:
Pull Request resolved: #4840

X-link: facebookresearch/FBGEMM#1865

Enable deterministic mode in Cutlass Blackwell Attention.  Previously,
all parts of the kernels were deterministic except for the dQ
computation in the backward pass, which performed reduction across
thread blocks using atomic add. This diff serializes the dQ
aggregation across thread blocks, ensuring full determinism throughout
the kernel execution.

Changes include:
- Rebased D79426672 onto master
- Enabled deterministic mode for sliding window attention
- Enabled deterministic mode for GQA in backward
- Added unit tests

Reviewed By: y-sq, jianyuh

Differential Revision: D81980925

fbshipit-source-id: 1fac8921d166e56692111c8a844a559c6036c7e2

Author: sryap

fbgemm_gpu/experimental/gen_ai/gen_ai/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py

@@ -104,7 +104,9 @@ def _cutlass_blackwell_fmha_backward(
     window_left: int = -1,
     window_right: int = -1,
     bottom_right: bool = True,
+    deterministic: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    deterministic = deterministic or torch.are_deterministic_algorithms_enabled()
     dout = maybe_contiguous(dout)
     q = maybe_contiguous(q)
     k = maybe_contiguous(k)
@@ -125,6 +127,7 @@ def _cutlass_blackwell_fmha_backward(
         window_size_left=window_left,
         window_size_right=window_right,
         bottom_right=bottom_right,
+        deterministic=deterministic,
     )
 
 
@@ -172,6 +175,7 @@ def forward(  # type: ignore
         seqlen_kv: Optional[torch.Tensor] = None,
         window_size: Tuple[int, int] = (-1, -1),
         bottom_right: bool = True,
+        deterministic: bool = False,
     ) -> torch.Tensor:
         # Check if this is generation phase (sq = 1)
         sq = q.shape[1]
@@ -232,6 +236,7 @@ def forward(  # type: ignore
             ctx.cu_seqlens_k = cu_seqlens_k
             ctx.is_gen = False
             ctx.bottom_right = bottom_right
+            ctx.deterministic = deterministic
             return out
 
     @staticmethod
@@ -248,6 +253,7 @@ def backward(ctx, dout: torch.Tensor, *args: Any) -> Tuple[  # type: ignore
         None,
         None,
         None,
+        None,
     ]:
         if ctx.is_gen:
             # For gen case, no backward pass is needed (generation is inference only)
@@ -272,8 +278,9 @@ def backward(ctx, dout: torch.Tensor, *args: Any) -> Tuple[  # type: ignore
             window_left,
             window_right,
             bottom_right=ctx.bottom_right,
+            deterministic=ctx.deterministic,
         )
-        return dq, dk, dv, None, None, None, None, None, None, None, None, None
+        return dq, dk, dv, None, None, None, None, None, None, None, None, None, None
 
 
 def cutlass_blackwell_fmha_func(
@@ -289,6 +296,7 @@ def cutlass_blackwell_fmha_func(
     seqlen_kv: torch.Tensor | None = None,
     window_size: tuple[int, int] | None = (-1, -1),
     bottom_right: bool = True,
+    deterministic: bool = False,
 ):
     return CutlassBlackwellFmhaFunc.apply(
         q,
@@ -303,4 +311,5 @@ def cutlass_blackwell_fmha_func(
         seqlen_kv,
         window_size,
         bottom_right,
+        deterministic,
     )

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

@@ -6,6 +6,7 @@ template <
     typename Element,
     typename ActiveMask,
     bool kIsVarlen,
+    bool kIsDeterministic,
     class... KernelOptions>
 std::tuple<at::Tensor, at::Tensor, at::Tensor> fmha_bwd(
     const at::Tensor& dO,
@@ -36,7 +37,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> fmha_bwd(
   using TileShape = Shape<_128, _128, _128>;
 
   using Operation = cutlass::fmha::device::
-      Sm100FmhaBwd<ProblemShapeType, Element, ElementAccumulator, TileShape, /*kIsMla=*/false, ActiveMask>;
+      Sm100FmhaBwd<ProblemShapeType, Element, ElementAccumulator, TileShape, /*kIsMla=*/false, ActiveMask, kIsDeterministic>;
 
   using StrideQ = Stride<int, _1, Stride<Stride<int, int>, int>>; // Q D    ((H_R, H_K), B)
   using StrideK = Stride<int, _1, Stride<Stride<_0, int>, int>>;  // K D    ((H_R, H_K), B)
@@ -219,6 +220,19 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> fmha_bwd(
       cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
           hw_info.device_id);
 
+  auto seqlen_q = kIsVarlen ? max_seq_len_q.value() : q.size(1);
+
+  int* dq_semaphore_ptr = nullptr;
+  at::Tensor dq_semaphore;
+  if (kIsDeterministic) {
+    auto kBlockM = cute::get<0>(TileShape{});
+    auto opts = q.options();
+    dq_semaphore = torch::zeros(
+        {(seqlen_q + kBlockM - 1) / kBlockM, B, H_Q},
+        opts.dtype(torch::kInt32));
+    dq_semaphore_ptr = static_cast<int*>(dq_semaphore.data_ptr());
+  }
+
   typename Operation::Arguments arguments{
     problem_shape,
     static_cast<Element*>(q.data_ptr()),
@@ -240,6 +254,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> fmha_bwd(
     static_cast<Element*>(dV.data_ptr()),
     stride_dV,
     softmax_scale,
+    dq_semaphore_ptr,
     window_size_left,
     window_size_right,
     hw_info};
@@ -264,7 +279,8 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dispatch_fmha_bwd(
     bool causal,
     int64_t window_size_left,
     int64_t window_size_right,
-    bool bottom_right
+    bool bottom_right,
+    bool deterministic
 ) {
   // This workaround initializes the CUDA context to prevent the 201 error
   // (invalid context).  When this function is invoked through PyTorch
@@ -294,11 +310,18 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dispatch_fmha_bwd(
   }
 
   auto dispatch_fmha =
-    [&](auto element, auto element_out, auto varlen, auto mask, auto... kernel_options) {
+    [&](
+      auto element,
+      auto element_out,
+      auto varlen,
+      auto deterministic,
+      auto mask,
+      auto... kernel_options) {
       return fmha_bwd<
         decltype(element),
         decltype(mask),
         varlen,
+        deterministic,
         decltype(kernel_options)...>
       (
         dOutput,
@@ -315,53 +338,69 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dispatch_fmha_bwd(
         window_size_right);
     };
 
-  auto dispatch_type = [&](auto varlen, auto mask) {
+  auto dispatch_type = [&](auto varlen, auto deterministic, auto mask) {
     if (query.dtype() == torch::kFloat16) {
-      return dispatch_fmha(cutlass::half_t{}, cutlass::half_t{}, varlen, mask);
+      return dispatch_fmha(
+          cutlass::half_t{}, cutlass::half_t{}, varlen, deterministic, mask);
     }
     else if (query.dtype() == torch::kBFloat16) {
       return dispatch_fmha(
-          cutlass::bfloat16_t{}, cutlass::bfloat16_t{}, varlen, mask);
+          cutlass::bfloat16_t{}, cutlass::bfloat16_t{}, varlen, deterministic, mask);
     }
     else if (query.dtype() == torch::kFloat8_e4m3fn) {
       return dispatch_fmha(
-          cutlass::float_e4m3_t{}, cutlass::bfloat16_t{}, varlen, mask);
+          cutlass::float_e4m3_t{}, cutlass::bfloat16_t{}, varlen, deterministic, mask);
     }
     TORCH_CHECK(false, "Unsupported dtype for q: ", query.dtype());
   };
 
-  auto dispatch_mask = [&](auto varlen) {
+  auto dispatch_mask = [&](auto varlen, auto deterministic) {
     if (causal) {
       if (bottom_right) {
-        return dispatch_type(varlen, CausalForBackwardMask</*kIsQBegin=*/false>{});
+        return dispatch_type(
+            varlen, deterministic, CausalForBackwardMask</*kIsQBegin=*/false>{});
       }
       else {
-        return dispatch_type(varlen, CausalForBackwardMask</*kIsQBegin=*/true>{});
+        return dispatch_type(
+            varlen, deterministic, CausalForBackwardMask</*kIsQBegin=*/true>{});
       }
     }
     else if (local) {
       if (bottom_right) {
-        return dispatch_type(varlen, LocalMaskForBackward</*kIsQBegin=*/false>{});
+        return dispatch_type(
+            varlen, deterministic, LocalMaskForBackward</*kIsQBegin=*/false>{});
       }
       else {
-        return dispatch_type(varlen, LocalMaskForBackward</*kIsQBegin=*/true>{});
+        return dispatch_type(
+            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(varlen, ResidualMaskForBackward{});
+      return dispatch_type(
+          varlen, deterministic, ResidualMaskForBackward{});
+    }
+    else {
+      return dispatch_type(
+          varlen, deterministic, NoMask{});
+    }
+  };
+
+  auto dispatch_deterministic = [&](auto varlen) {
+    if (deterministic) {
+      return dispatch_mask(varlen, std::bool_constant<true>{});
     }
     else {
-      return dispatch_type(varlen, NoMask{});
+      return dispatch_mask(varlen, std::bool_constant<false>{});
     }
   };
 
   if (max_seq_len_q.has_value()) {
-    return dispatch_mask(std::bool_constant<true>{});
+    return dispatch_deterministic(std::bool_constant<true>{});
   } else {
     TORCH_CHECK(query.dim() == 4, "q must be [B, M, H, D] for fixed length")
-    return dispatch_mask(std::bool_constant<false>{});
+    return dispatch_deterministic(std::bool_constant<false>{});
   }
 }
 
@@ -383,7 +422,8 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
         "    bool causal=False, "
         "    int window_size_left=-1, "
         "    int window_size_right=-1, "
-        "    bool bottom_right=True"
+        "    bool bottom_right=True, "
+        "    bool deterministic=False"
         ") -> (Tensor, Tensor, Tensor)"
   );
 }

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/device/fmha_device_bwd.hpp

@@ -58,7 +58,8 @@ template<
     class ElementAccumulator,
     class TileShape,
     bool IsMla,
-    class Mask
+    class Mask,
+    bool IsDeterministic=false
 >
 class Sm100FmhaBwd {
 private:
@@ -123,6 +124,8 @@ class Sm100FmhaBwd {
 
     ElementAccumulator softmax_scale;
 
+    int* ptr_dq_semaphore;
+
     int window_size_left = -1;
     int window_size_right = -1;
 
@@ -138,7 +141,7 @@ class Sm100FmhaBwd {
 
   using OperationMha = cutlass::fmha::device::FMHA<
       cutlass::fmha::kernel::Sm100FmhaBwdKernelTmaWarpSpecialized<
-          ProblemShape, Element, ElementAccumulator, TileShape, Mask
+          ProblemShape, Element, ElementAccumulator, TileShape, Mask, IsDeterministic
       >
   >;
 
@@ -223,7 +226,10 @@ class Sm100FmhaBwd {
         scaled_lse, to_bwd_stride(stride_scaled_lse),
         sum_OdO, to_bwd_stride(stride_sum_OdO),
         dQ_acc, to_bwd_stride(stride_dQ),
-        args.softmax_scale, args.window_size_left, args.window_size_right},
+        args.softmax_scale,
+        args.ptr_dq_semaphore,
+        args.window_size_left,
+        args.window_size_right },
       { args.ptr_dK, to_bwd_stride(args.stride_dK),
         args.ptr_dV, to_bwd_stride(args.stride_dV) },
       args.hw_info