Support batch size > 1 when enable CP (#23269)

Co-authored-by: Shunkang <182541032+Shunkangz@users.noreply.github.co>
Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>

test/registered/kernels/test_cp_prefix_len_fa3_parity.py

@@ -77,7 +77,7 @@ class TestCPPrefixLenFA3Parity(CustomTestCase):
 
         def _call_meta(rank: int):
             return prepare_context_parallel_metadata(
-                padded_extend, rank, cp_size, seqs_len, extend_lens=extend_lens
+                padded_extend, rank, cp_size, seqs_len, extend_seqs_len=extend_lens
             )
 
         # Exercise the non-DSA branch; the DSA branch uses a separate

test/registered/kernels/test_mla_cp_fa3_parity.py

@@ -126,8 +126,15 @@ def _cp_attn_for_rank(
         kv_len_next_tensor=torch.tensor(
             [(b_next + 1) * block_size], dtype=torch.int32, device=DEVICE
         ),
-        actual_seq_q_prev=block_size,
-        actual_seq_q_next=block_size,
+        cu_seqlens_q_prev_tensor=torch.tensor(
+            [0, block_size], dtype=torch.int32, device=DEVICE
+        ),
+        cu_seqlens_q_next_tensor=torch.tensor(
+            [0, block_size], dtype=torch.int32, device=DEVICE
+        ),
+        max_seqlen_q_prev=block_size,
+        max_seqlen_q_next=block_size,
+        total_q_prev_tokens=block_size,
     )
     fb = SimpleNamespace(attn_cp_metadata=cp_meta)