add assertion on len == 1

Author: chunyuan-w

torch/_inductor/kernel/flex_attention.py

@@ -814,6 +814,8 @@ def flex_attention(
         choices: List[Any] = []
         input_nodes = [query, key, value, kv_indices]
         if score_mod_other_buffers and mask_mod_other_buffers:
+            assert len(score_mod_other_buffers) == 1
+            assert len(mask_mod_other_buffers) == 1
             input_nodes += [score_mod_other_buffers[0], mask_mod_other_buffers[0]]        
         CppMHATemplate.add_choices(
             choices=choices,