Use int64_t instead of int32_t in decode attention kernel to avoid index calc overflow (#5095)

Summary:
Pull Request resolved: #5095

X-link: https://github.com/facebookresearch/FBGEMM/pull/2101

For KV cache size >4GB, int32_t is not enough to calculate the index.

Reviewed By: Aya-ZIbra

Differential Revision: D85631787

fbshipit-source-id: a9695be5a6d79b9abd3a349f780878d766ac2ceb

Author: Alkaid-Benetnash

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/77_blackwell_fmha_gen.cu

@@ -335,7 +335,7 @@ struct ExampleRunner {
 
   using StrideQ = Stride<_0, _1, Stride<Stride<int, int>, int>>;
   using StrideNewK = Stride<_0, _1, Stride<Stride<_0, int>, int>>;
-  using StrideCacheK = Stride<int, _1, Stride<Stride<_0, int>, int>>;
+  using StrideCacheK = Stride<int, _1, Stride<Stride<_0, int64_t>, int64_t>>;
   using StrideNewV = StrideNewK;
   using StrideCacheV = StrideCacheK;
   using StrideO = StrideQ;
@@ -474,7 +474,7 @@ struct ExampleRunner {
 
     stride_q = make_stride(_0{}, _1{}, make_stride(make_stride(options.d, options.d * size<3,0,0>(result)), options.d * size<3,0>(result)));
     stride_new_k = make_stride(_0{}, _1{}, make_stride(make_stride(_0{}, options.d), options.d * size<3,0,1>(result)));
-    stride_cache_k = make_stride(options.d * size<3,0,1>(result), _1{}, make_stride(make_stride(_0{}, options.d), options.d * size<3,0,1>(result) * get<1>(result)));
+    stride_cache_k = make_stride(options.d * size<3,0,1>(result), _1{}, make_stride(make_stride(_0{}, static_cast<int64_t>(options.d)), static_cast<int64_t>(options.d) * size<3,0,1>(result) * get<1>(result)));
 
     stride_new_v = stride_new_k;
     stride_cache_v = stride_cache_k;

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

@@ -93,9 +93,9 @@ struct GenRunner {
 
   using StrideQ =
       Stride<_0, _1, Stride<Stride<int, int>, int>>; // Q D ((H, Hr), B)
-  using StrideNewK = Stride<_0, _1, Stride<Stride<_0, int>, int>>;
+  using StrideNewK = Stride<_0, _1, Stride<Stride<_0, int64_t>, int64_t>>;
   using StrideCacheK =
-      Stride<int, _1, Stride<Stride<_0, int>, int>>; // K D ((H, Hr), B)
+      Stride<int, _1, Stride<Stride<_0, int64_t>, int64_t>>; // K D ((H, Hr), B)
   using StrideNewV = StrideNewK;
   using StrideCacheV = StrideCacheK;
   using StrideO = StrideQ;
@@ -198,13 +198,13 @@ struct GenRunner {
     stride_new_k = make_stride(
         _0{},
         _1{},
-        make_stride(make_stride(_0{}, options.d), options.d * options.h_k));
+        make_stride(make_stride(_0{}, static_cast<int64_t>(options.d)), static_cast<int64_t>(options.d * options.h_k)));
     stride_cache_k = make_stride(
         options.d * options.h_k,
         _1{},
         make_stride(
-            make_stride(_0{}, options.d),
-            options.d * options.h_k * options.sk));
+            make_stride(_0{}, static_cast<int64_t>(options.d)),
+            static_cast<int64_t>(options.d * options.h_k * options.sk)));
 
     stride_new_v = stride_new_k;
     stride_cache_v = stride_cache_k;