Add tests for non divisible inputs for flex decoding (pytorch#143214)

Pull Request resolved: pytorch#143214
Approved by: https://github.com/drisspg

Author: joydddd

test/inductor/test_flex_decoding.py

@@ -281,7 +281,7 @@ def _check_out(
 
     def run_test(
         self,
-        score_mod: Optional[Callable],
+        score_mod: Optional[Callable] = None,
         dtype: torch.dtype = torch.float16,
         Q_B: int = B,
         Q_H: int = Hq,
@@ -348,7 +348,7 @@ def run_test_with_call(
         if not golden_call:
             golden_call = sdpa_call
         q = torch.randn(
-            (Q_B, KV_H, Q_S * (Q_H // KV_H), Q_D),
+            (Q_B, KV_H, Q_S, Q_D),
             dtype=dtype,
             device="cuda",
             requires_grad=False,
@@ -850,6 +850,61 @@ def seq_mask_mod(score, b, h, q, kv):
         self.run_test(seq_mask_mod, dtype)
         self.run_test_with_paged_attention(seq_mask_mod, dtype)
 
+    @supported_platform
+    def test_non_divisible_offset_mask(self):
+        KV_S = S - 3
+        offset_tensor = torch.tensor(S // 2 - 3, device="cuda", dtype=torch.int32)
+
+        def mask_mod(b, h, q, kv):
+            return kv >= q + offset_tensor
+
+        block_mask = create_block_mask(mask_mod, B, 1, 1, KV_S)
+        self.run_test(KV_S=KV_S, block_mask=block_mask)
+
+    @supported_platform
+    def test_non_divisible_offset_mask_with_captured_buffer(self):
+        KV_S = S - 3
+        offset_kv = torch.randn(KV_S, device="cuda", dtype=torch.bfloat16)
+        offset_tensor = torch.tensor(S // 2 - 3, device="cuda", dtype=torch.int32)
+
+        def score_mod(score, b, h, q, kv):
+            return score + offset_kv[kv]
+
+        def mask_mod(b, h, q, kv):
+            return kv >= q + offset_tensor
+
+        block_mask = create_block_mask(mask_mod, B, 1, 1, KV_S)
+        self.run_test(KV_S=KV_S, block_mask=block_mask, score_mod=score_mod)
+
+    @supported_platform
+    def test_non_divisible_multi_token_offset_mask(self):
+        KV_S = S - 3
+        Q_S = 3
+        offset_tensor = torch.tensor(S // 2 - 1, device="cuda", dtype=torch.int32)
+
+        def mask_mod(b, h, q, kv):
+            return kv >= q + offset_tensor
+
+        block_mask = create_block_mask(mask_mod, B, 1, Q_S, KV_S)
+        self.run_test(Q_S=Q_S, KV_S=KV_S, block_mask=block_mask)
+
+    @supported_platform
+    def test_non_divisible_multi_token_offset_mask_with_captured_buffer(self):
+        KV_S = S - 3
+        Q_S = 3
+        offset_kv = torch.randn(KV_S, device="cuda", dtype=torch.bfloat16)
+        offset_q = torch.randn(Q_S, device="cuda", dtype=torch.bfloat16)
+        offset_tensor = torch.tensor(S // 2 - 3, device="cuda", dtype=torch.int32)
+
+        def score_mod(score, b, h, q, kv):
+            return score + offset_kv[kv] + offset_q[q]
+
+        def mask_mod(b, h, q, kv):
+            return kv >= q + offset_tensor
+
+        block_mask = create_block_mask(mask_mod, B, 1, Q_S, KV_S)
+        self.run_test(Q_S=Q_S, KV_S=KV_S, block_mask=block_mask, score_mod=score_mod)
+
     @supported_platform
     @common_utils.parametrize("dtype", test_dtypes_fast)
     def test_load_from_bias_seq_only(self, dtype):