Fix overflow in attention workspace calculation (pytorch#4939)

Summary:
X-link: facebookresearch/FBGEMM#1961

Pull Request resolved: pytorch#4939

During workspace calculation some of the additions risked overflowing int32 and returning an invalid value. This diff makes sure that we're doing higher precision arithmetic in the workspace to prevent this issue.

Reviewed By: q10

Differential Revision: D83370379

fbshipit-source-id: b481d16ce4053a21ac19f1060eec6d57bbeb9eda

Author: jwfromm

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

@@ -267,14 +267,15 @@ class Sm100FmhaBwd {
     auto [Q_, K, D, D_VO, HB] = args.problem_shape;
     auto [H, B] = product_each(HB);
     D = cutlass::round_up(D, 8);  // Alignment
-    int Q = cutlass::round_up(static_cast<int>(Q_), 8);  // Alignment
+    size_t Q = cutlass::round_up(static_cast<int>(Q_), 8);  // Alignment
     size_t workspace_bytes = 0;
+    size_t accum_size = sizeof(ElementAccumulator);
     // OdO vector
-    workspace_bytes += B*H*Q * sizeof(ElementAccumulator);
+    workspace_bytes += static_cast<size_t>(B)*static_cast<size_t>(H)*Q * accum_size;
     // scaled LSE vector
-    workspace_bytes += B*H*Q * sizeof(ElementAccumulator);
+    workspace_bytes += static_cast<size_t>(B)*static_cast<size_t>(H)*Q * accum_size;
     // FP32 versions of outputs that are churned (start off with Q only)
-    workspace_bytes += B*H*Q*D * sizeof(ElementAccumulator);
+    workspace_bytes += static_cast<size_t>(B)*static_cast<size_t>(H)*Q*static_cast<size_t>(D) * accum_size;
     return workspace_bytes;
   }