解决大batch长seq崩溃问题

csrc/attention/attention_kernels_opt_tc.cu

@@ -1054,9 +1054,9 @@ void paged_attention_v2_launcher_opt_tc(
   static float* max_logits_ptr = nullptr;
   static T* tmp_out_ptr = nullptr;
   if(exp_sums_ptr == nullptr){
-      hipMalloc(&exp_sums_ptr, 10000000); // 10m
-      hipMalloc(&max_logits_ptr, 10000000); // 10m
-      hipMalloc(&tmp_out_ptr, 400000000); // 400m
+      hipMalloc(&exp_sums_ptr, 1000000); // 1m
+      hipMalloc(&max_logits_ptr, 1000000); // 1m
+      hipMalloc(&tmp_out_ptr, 100000000); // 100m
   }
   const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

csrc/attention/attention_with_mask_kernels_opt_tc.cu

@@ -959,9 +959,9 @@ void paged_attention_v2_launcher_opt_tc_with_mask(
   static float* max_logits_ptr = nullptr;
   static T* tmp_out_ptr = nullptr;
   if(exp_sums_ptr == nullptr){
-      hipMalloc(&exp_sums_ptr, 10000000); // 10m
-      hipMalloc(&max_logits_ptr, 10000000); // 10m
-      hipMalloc(&tmp_out_ptr, 400000000); // 400m
+      hipMalloc(&exp_sums_ptr, 1000000); // 1m
+      hipMalloc(&max_logits_ptr, 1000000); // 1m
+      hipMalloc(&tmp_out_ptr, 100000000); // 100m
   }
   const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();