Make test changes for from_segmet_ids_and_pos API change. Also some nits.

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Author: KshitijLakhani

tests/jax/test_fused_attn.py

@@ -547,13 +547,14 @@ def _setup_inputs(self):
         else:
             self.softmax_offset = None
 
-        def gen_valid(bs, max_seqlen, pad_ratio):
+        def generate_valid_segment_ids_and_pos(bs, max_seqlen, pad_ratio):
             pad_len = int(max_seqlen * pad_ratio)
             valid_len = max_seqlen - pad_len
             tokens = jnp.concatenate([jnp.ones((bs, valid_len)), jnp.zeros((bs, pad_len))], axis=-1)
-            return tokens, jnp.logical_not(tokens)
+            segment_pos = jnp.broadcast_to(jnp.arange(max_seqlen, dtype=tokens.dtype), tokens.shape)
+            return tokens, segment_pos, jnp.logical_not(tokens)
 
-        def generate_random_segment_ids(
+        def generate_random_segment_ids_and_pos(
             batch_size,
             sequence_length,
             num_segments,
@@ -601,7 +602,7 @@ def generate_random_segment_ids(
             return segment_ids, segment_pos, segment_pad
 
         if self.qkv_layout.is_thd():
-            self.segment_ids_q, self.segment_pos_q, self.pad_q = generate_random_segment_ids(
+            self.segment_ids_q, self.segment_pos_q, self.pad_q = generate_random_segment_ids_and_pos(
                 self.batch_size, self.max_seqlen_q, self.num_segments_per_seq, seed=42
             )
             self.seqlens_q, self.offsets_q = get_seqlens_and_offsets(self.segment_ids_q)
@@ -617,7 +618,7 @@ def generate_random_segment_ids(
                     self.window_size is not None or self.attn_mask_type.is_bottom_right()
                 ):  # SWA or BRCM requires kv_len >= q_len
                     min_segment_len = self.seqlens_q
-                self.segment_ids_kv, self.segment_pos_kv, self.pad_kv = generate_random_segment_ids(
+                self.segment_ids_kv, self.segment_pos_kv, self.pad_kv = generate_random_segment_ids_and_pos(
                     self.batch_size,
                     self.max_seqlen_kv,
                     self.num_segments_per_seq,
@@ -626,13 +627,12 @@ def generate_random_segment_ids(
                 )
             self.seqlens_kv, self.offsets_kv = get_seqlens_and_offsets(self.segment_ids_kv)
         else:
-            self.segment_ids_q, self.pad_q = gen_valid(
+            self.segment_ids_q, self.segment_pos_q, self.pad_q = generate_valid_segment_ids_and_pos(
                 self.batch_size, self.max_seqlen_q, pad_ratio
             )
-            self.segment_ids_kv, self.pad_kv = gen_valid(
+            self.segment_ids_kv, self.segment_pos_kv, self.pad_kv = generate_valid_segment_ids_and_pos(
                 self.batch_size, self.max_seqlen_kv, pad_ratio
             )
-            self.segment_pos_q = self.segment_pos_kv = None
             self.seqlens_q = self.seqlens_kv = self.offsets_q = self.offsets_kv = None
 
         # For reference code
@@ -682,24 +682,15 @@ def generate_random_segment_ids(
                         (self.offsets_q, self.offsets_kv),
                     )
                 case SeqDescFormat.SegmentIDs:
-                    # Exercise the path to generate the segment_pos in from_segment_ids_and_pos()
-                    # if no CP and load balancing, else explicitly pass the segment_pos
+                    # from_segment_ids_and_pos requires explicit segment_pos.
                     self.sequence_desciptor = SequenceDescriptor.from_segment_ids_and_pos(
                         (
                             self.cp_reorder_fn(self.segment_ids_q),
                             self.cp_reorder_fn(self.segment_ids_kv),
                         ),
                         (
-                            (
-                                self.cp_reorder_fn(self.segment_pos_q),
-                                self.cp_reorder_fn(self.segment_pos_kv),
-                            )
-                            if self.cp_size > 1 and self.cp_load_balanced
-                            else None
-                        ),
-                        is_thd=self.qkv_layout.is_thd(),
-                        is_segment_ids_reordered=(
-                            True if self.cp_size > 1 and self.cp_load_balanced else False
+                            self.cp_reorder_fn(self.segment_pos_q),
+                            self.cp_reorder_fn(self.segment_pos_kv),
                         ),
                     )
                 case _:
@@ -727,9 +718,7 @@ def generate_random_segment_ids(
                 case SeqDescFormat.SegmentIDs:
                     self.sequence_desciptor = SequenceDescriptor.from_segment_ids_and_pos(
                         (self.segment_ids_q, self.segment_ids_kv),
-                        None,
-                        is_thd=self.qkv_layout.is_thd(),
-                        is_segment_ids_reordered=False,
+                        (self.segment_pos_q, self.segment_pos_kv),
                     )
                 case _:
                     raise ValueError(f"Unknown {self.seq_desc_format=}")