[JAX] Fix incorrect calculation of segment pos from segment ids in user-facing API

tests/jax/test_fused_attn.py

@@ -677,6 +677,8 @@ def generate_random_segment_ids(
                             self.cp_reorder_fn(self.segment_pos_q),
                             self.cp_reorder_fn(self.segment_pos_kv),
                         ),
+                        self.qkv_layout.is_thd(),
+                        self.cp_size > 1 and self.cp_load_balanced,
                     )
                 case _:
                     raise ValueError(f"Unknown {self.seq_desc_format=}")
@@ -704,6 +706,8 @@ def generate_random_segment_ids(
                     self.sequence_desciptor = SequenceDescriptor.from_segment_ids_and_pos(
                         (self.segment_ids_q, self.segment_ids_kv),
                         None,
+                        self.qkv_layout.is_thd(),
+                        self.cp_size > 1 and self.cp_load_balanced,
                     )
                 case _:
                     raise ValueError(f"Unknown {self.seq_desc_format=}")

transformer_engine/jax/attention.py

@@ -791,14 +791,18 @@ def from_seqlens_and_offsets(
         q_offsets, kv_offsets = cls._expand_to_pair(seq_offsets)
         return cls(seqlens=(q_seqlens, kv_seqlens), seq_offsets=(q_offsets, kv_offsets))
 
+    # TODO(KshitijLakhani), TODO(mgoldfarb-nvidia): Consider adding support for THD layout (non load balanced).
     @classmethod
     def from_segment_ids_and_pos(
         cls,
         segment_ids: Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]],
         segment_pos: Optional[Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]]] = None,
+        is_thd: bool = False,
+        is_load_balanced: bool = False,
     ) -> SequenceDescriptor:
         """
-        Experimental factory method for inputs with segment IDs and optional positions. (THD)
+        Experimental factory method for inputs with segment IDs and optional positions.
+        segment_pos = None to be used only for : BSHD without load balancing
         Args:
             segment_ids(Tuple(jnp.ndarray, jnp.ndarray)) = (q_segment_ids, kv_segment_ids):
                 - q_segment_ids (jnp.ndarray):
@@ -812,14 +816,25 @@ def from_segment_ids_and_pos(
                   The position inside each segment for query, with shape [batch, max_seqlen].
                 - kv_segment_pos (jnp.ndarray):
                   The position inside each segment for key, value, with shape [batch, max_seqlen].
+            is_thd(bool): If True, QKVLayout is of type THD, else it is BSHD
+            is_load_balanced(bool): If True, CP is being used and the inputs have been load balanced.
         Return:
             A SequenceDescriptor with segment_ids/segment_pos initialized.
         """
         q_seg_ids, kv_seg_ids = cls._expand_to_pair(segment_ids)
 
-        if segment_pos is not None:
-            segment_pos = cls._expand_to_pair(segment_pos)
-        else:
+        # If using defaults
+        if segment_pos is None:
+            # Segment pos is not calculated implicitly for THD cases and Load balancing cases
+            assert (
+                not is_load_balanced
+            ), f"segment_pos = None default arg is not supported for load balanced inputs"
+            assert not is_thd, f"segment_pos = None default arg is not supported for THD layouts"
+            warnings.warn(
+                "segment_pos = None is only acceptable if using BSHD and no load balancing. For all"
+                " other cases, segment_pos must be passed explicitly",
+                UserWarning,
+            )
 
             def generate_default_pos(segment_ids):
                 seqlen = segment_ids.shape[-1]
@@ -828,6 +843,8 @@ def generate_default_pos(segment_ids):
             q_seg_pos = generate_default_pos(q_seg_ids)
             kv_seg_pos = generate_default_pos(kv_seg_ids)
             segment_pos = (q_seg_pos, kv_seg_pos)
+        else:  # Explicitly passed segment_pos
+            segment_pos = cls._expand_to_pair(segment_pos)
 
         return cls(
             segment_ids=(q_seg_ids, kv_seg_ids),