[Pallas] Improve segment_ids API UX (#7037)

Summary: Adds comment for the shape of the segment_ids. Test Plan: Skip CI.
pytorch · May 8, 2024 · 5bbe5c8 · 5bbe5c8
1 parent 887d344
commit 5bbe5c8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_xla/experimental/custom_kernel.py b/torch_xla/experimental/custom_kernel.py
@@ -421,8 +421,8 @@ def flash_attention(
     k,  # [batch_size, num_heads, kv_seq_len, d_model]
     v,  # [batch_size, num_heads, kv_seq_len, d_model]
     causal=False,
-    q_segment_ids=None,
-    kv_segment_ids=None,
+    q_segment_ids=None,  # [batch_size, q_seq_len]
+    kv_segment_ids=None,  # [batch_size, kv_seq_len]
     sm_scale=1.0,
     *,
     partition_spec=None,