Fix, or rather "port", bug fix for sdpa

Differential Revision: D73640471

Pull Request resolved: #10466

Author: kimishpatel

extension/llm/custom_ops/op_sdpa_impl.h

@@ -968,27 +968,36 @@ void cpu_flash_attention(
                 tmp_max);
           }
           tmp_max = qk_max_data[row] > tmp_max ? qk_max_data[row] : tmp_max;
-          // qk <- exp(qk - max) and sum per row
-          tmp_sum = tmp_max;
-          _exp_reduce_sum_fusion_kernel(
-              qk_data + row * kvBlockSize,
-              kvBlockSize,
-              conditional_data_ptr(qk_data, qk_reduced_data) +
-                  row * kvBlockSize,
-              tmp_sum);
-          // exp_tmp <- exp(max[row] - max)
-          exp_tmp = std::exp(qk_max_data[row] - tmp_max);
-          // sum[row] <- sum + exp_tmp * sum[row]
-          qk_sum_data[row] = tmp_sum + exp_tmp * qk_sum_data[row];
-          // max[row] <- max
-          qk_max_data[row] = tmp_max;
-          // dst <- dst * exp_tmp
-          if (n > 0) {
-            vec::map<accum_t>(
-                [exp_tmp](Vec x) { return x * Vec(exp_tmp); },
-                dst_data + row * headSize,
-                dst_data + row * headSize,
-                headSize);
+          if (tmp_max == -std::numeric_limits<accum_t>::infinity()) {
+            // to avoid `nan = exp2f(-inf - (-inf))`
+            fill_stub(
+                conditional_data_ptr(qk_data, qk_reduced_data) +
+                    row * kvBlockSize,
+                static_cast<scalar_t>(0),
+                kvBlockSize);
+          } else {
+            // qk <- exp(qk - max) and sum per row
+            tmp_sum = tmp_max;
+            _exp_reduce_sum_fusion_kernel(
+                qk_data + row * kvBlockSize,
+                kvBlockSize,
+                conditional_data_ptr(qk_data, qk_reduced_data) +
+                    row * kvBlockSize,
+                tmp_sum);
+            // exp_tmp <- exp(max[row] - max)
+            exp_tmp = std::exp(qk_max_data[row] - tmp_max);
+            // sum[row] <- sum + exp_tmp * sum[row]
+            qk_sum_data[row] = tmp_sum + exp_tmp * qk_sum_data[row];
+            // max[row] <- max
+            qk_max_data[row] = tmp_max;
+            // dst <- dst * exp_tmp
+            if (n > 0) {
+              vec::map<accum_t>(
+                  [exp_tmp](Vec x) { return x * Vec(exp_tmp); },
+                  dst_data + row * headSize,
+                  dst_data + row * headSize,
+                  headSize);
+            }
           }
         }
 

extension/llm/custom_ops/test_sdpa_with_kv_cache.py

@@ -53,12 +53,13 @@ def setUp(self):
         self.mask = torch.triu(self.mask, diagonal=1)
         self.use_mask_with_custom_op = False
         self.is_causal = False
+        self.start_pos = 0
 
     def test_sdpa_with_cache_no_mqa_1(self):
         q = torch.rand((1, 1, 8, 4))
         k = torch.rand((1, 1, 8, 4))
         v = torch.rand((1, 1, 8, 4))
-        start_pos = 0
+        start_pos = self.start_pos
         seq_len = q.size(1)
         attn_mask = self.mask[start_pos : start_pos + seq_len, :]
         attn_mask = attn_mask[:, : start_pos + seq_len]
@@ -238,6 +239,38 @@ def setUp(self):
         self.use_mask_with_custom_op = True
 
 
+class SDPAWithAttentionMaskLongSequenceTest(SDPATest):
+
+    def setUp(self):
+        SDPATest.setUp(self)
+        max_context_len = 700
+        context_window_len = 60
+        self.k_cache = torch.zeros((1, 700, 8, 4))
+        self.v_cache = torch.zeros((1, 700, 8, 4))
+        causal_mask = torch.tril(
+            torch.ones(
+                max_context_len,
+                max_context_len,
+                dtype=torch.bool,
+                device="cpu",
+            )
+        )
+        causal_mask2 = torch.tril(
+            torch.ones(
+                max_context_len,
+                max_context_len,
+                dtype=torch.bool,
+                device="cpu",
+            ),
+            diagonal=-context_window_len,
+        )
+        mask = torch.logical_xor(causal_mask, causal_mask2)
+        self.mask = torch.where(mask == True, 0.0, float("-inf"))  # noqa: E712
+
+        self.use_mask_with_custom_op = True
+        self.start_pos = 575
+
+
 class SDPAWithCausalTest(SDPATest):
 
     def setUp(self):