better error handling when window size is -1 and no max seq len is provided (#4979)

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

Pull Request resolved: #4979

Improve error handling when one of the window size is not provided and we don't know what max seq len is.

Reviewed By: Aya-ZIbra

Differential Revision: D84021746

fbshipit-source-id: f0142f27b5e365a0f8698e29f9fe2f1f940cb859

Author: henrylhtsang

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

@@ -41,10 +41,16 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dispatch_fmha_bwd(
       causal = false; // Use local attention instead of causal
     }
     // Expand -1 window sizes to full sequence length if available
-    if (window_size_left < 0 && max_seq_len_k.has_value()) {
+    if (window_size_left < 0) {
+      TORCH_CHECK(
+          max_seq_len_k.has_value(),
+          "window_size_left is negative but max_seq_len_k is not provided");
       window_size_left = max_seq_len_k.value();
     }
-    if (window_size_right < 0 && max_seq_len_k.has_value()) {
+    if (window_size_right < 0) {
+      TORCH_CHECK(
+          max_seq_len_k.has_value(),
+          "window_size_right is negative but max_seq_len_k is not provided");
       window_size_right = max_seq_len_k.value();
     }
   }

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

@@ -26,10 +26,16 @@ std::tuple<at::Tensor, at::Tensor> dispatch_fmha_fwd(
       causal = false; // Use local attention instead of causal
     }
     // Expand -1 window sizes to full sequence length if available
-    if (window_size_left < 0 && max_seq_len_k.has_value()) {
+    if (window_size_left < 0) {
+      TORCH_CHECK(
+          max_seq_len_k.has_value(),
+          "window_size_left is negative but max_seq_len_k is not provided");
       window_size_left = max_seq_len_k.value();
     }
-    if (window_size_right < 0 && max_seq_len_k.has_value()) {
+    if (window_size_right < 0) {
+      TORCH_CHECK(
+          max_seq_len_k.has_value(),
+          "window_size_right is negative but max_seq_len_k is not provided");
       window_size_right = max_seq_len_k.value();
     }
   }