diff --git a/deepspeed/compile/custom_ops/__init__.py b/deepspeed/compile/custom_ops/__init__.py
index 5885328f1fd8..31050595141f 100644
--- a/deepspeed/compile/custom_ops/__init__.py
+++ b/deepspeed/compile/custom_ops/__init__.py
@@ -3,10 +3,11 @@
 
 # DeepSpeed Team
 
-from .all_to_all import all_to_all
+from .all_to_all import all_gather_sequence, all_to_all, aggregate_loss
 from .tp_collectives import copy_to_tp_region, gather_from_tp_region, reduce_from_tp_region
 from . import sp_dp_registry
 
 __all__ = [
-    "all_to_all", "copy_to_tp_region", "gather_from_tp_region", "reduce_from_tp_region", "sp_dp_registry", "sp_compat"
+    "all_gather_sequence", "all_to_all", "aggregate_loss", "copy_to_tp_region", "gather_from_tp_region",
+    "reduce_from_tp_region", "sp_dp_registry", "sp_compat"
 ]
diff --git a/deepspeed/compile/custom_ops/all_to_all.py b/deepspeed/compile/custom_ops/all_to_all.py
index 3307bbc527ff..9e99ea4f2164 100644
--- a/deepspeed/compile/custom_ops/all_to_all.py
+++ b/deepspeed/compile/custom_ops/all_to_all.py
@@ -3,6 +3,8 @@
 
 # DeepSpeed Team
 
+from typing import Tuple
+
 import torch
 import deepspeed.comm as dist
 from torch.utils._sympy.functions import FloorDiv
@@ -31,6 +33,9 @@ def all_to_all(
 
     if scatter_idx == 1:
         N, local_S = dim1, dim2
+        if N % sp_size() != 0:
+            raise ValueError(f"AutoSP requires the Q/K/V head count ({N}) to be divisible by "
+                             f"sequence_parallel_size ({sp_size()})")
         input_t = input.reshape(B, sp_size(), N // sp_size(), local_S, H)
         input_t = input_t.permute(1, 0, 2, 3, 4).contiguous()
 
@@ -90,3 +95,58 @@ def _all_to_all_backward(ctx, grad):
 
 
 torch.library.register_autograd("autosp::all_to_all", _all_to_all_backward, setup_context=_all_to_all_backward_setup)
+
+
+@torch.library.custom_op("autosp::all_gather_sequence", mutates_args=())
+def all_gather_sequence(input: torch.Tensor, dim: int) -> torch.Tensor:
+    """Gather a local attention-mask dimension across the current SP group."""
+    assert is_setup(), 'Incorrect initialization of SP/DP mesh.'
+    gid = dist.get_rank() // sp_size()
+    group = get_group(gid)
+    outputs = [torch.empty_like(input) for _ in range(sp_size())]
+    dist.all_gather(outputs, input, group=group)
+    return torch.cat(outputs, dim=dim)
+
+
+@torch.library.register_fake("autosp::all_gather_sequence")
+def all_gather_sequence_fake(input: torch.Tensor, dim: int):
+    output_shape = list(input.shape)
+    output_shape[dim] *= sp_size()
+    return input.new_empty(output_shape)
+
+
+@torch.library.custom_op("autosp::aggregate_loss", mutates_args=())
+def aggregate_loss(loss: torch.Tensor, valid_tokens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Return the valid-token-weighted mean loss on every SP rank."""
+    assert is_setup(), 'Incorrect initialization of SP/DP mesh.'
+    gid = dist.get_rank() // sp_size()
+    group = get_group(gid)
+    finite_loss = torch.where(valid_tokens > 0, loss, torch.zeros_like(loss))
+    total = finite_loss * valid_tokens.to(loss.dtype)
+    total_tokens = valid_tokens.clone()
+    dist.all_reduce(total, group=group)
+    dist.all_reduce(total_tokens, group=group)
+    weight = valid_tokens.to(loss.dtype) / total_tokens.clamp_min(1).to(loss.dtype)
+    return total / total_tokens.clamp_min(1).to(loss.dtype), weight
+
+
+@torch.library.register_fake("autosp::aggregate_loss")
+def aggregate_loss_fake(loss: torch.Tensor, valid_tokens: torch.Tensor):
+    return torch.empty_like(loss), torch.empty_like(loss)
+
+
+def _aggregate_loss_backward_setup(ctx, inputs, output):
+    _, weight = output
+    ctx.mark_non_differentiable(weight)
+    ctx.save_for_backward(weight)
+
+
+def _aggregate_loss_backward(ctx, grad_loss, grad_weight):
+    del grad_weight
+    (weight, ) = ctx.saved_tensors
+    return grad_loss * weight, None
+
+
+torch.library.register_autograd("autosp::aggregate_loss",
+                                _aggregate_loss_backward,
+                                setup_context=_aggregate_loss_backward_setup)
diff --git a/deepspeed/compile/custom_ops/sp_dp_registry.py b/deepspeed/compile/custom_ops/sp_dp_registry.py
index a93707032959..a0d0c4368dfc 100644
--- a/deepspeed/compile/custom_ops/sp_dp_registry.py
+++ b/deepspeed/compile/custom_ops/sp_dp_registry.py
@@ -50,7 +50,17 @@ def dp_size():
 def populate_registry(SP_SIZE, DP_SIZE):
     """ Populate rank to SP/DP mesh index.  """
 
+    world_size = dist.get_world_size()
+    if SP_SIZE * DP_SIZE != world_size:
+        raise ValueError(f"AutoSP mesh ({SP_SIZE} x {DP_SIZE}) must cover the distributed world size ({world_size})")
+
     if GROUP_REGISTRY.get('is_reg', False):
+        current_mesh = (GROUP_REGISTRY['SP_SIZE'], GROUP_REGISTRY['DP_SIZE'])
+        requested_mesh = (SP_SIZE, DP_SIZE)
+        if current_mesh != requested_mesh:
+            raise RuntimeError(f"AutoSP process groups are already initialized for mesh {current_mesh}, "
+                               f"but mesh {requested_mesh} was requested. Reinitialize the distributed "
+                               "process before changing sequence_parallel_size.")
         return
 
     group_listing = []
diff --git a/deepspeed/compile/fx.py b/deepspeed/compile/fx.py
index 232f967fa328..4ecc3c52bf73 100644
--- a/deepspeed/compile/fx.py
+++ b/deepspeed/compile/fx.py
@@ -172,7 +172,8 @@ def find_node_by_name(gm: GraphModule, name: str) -> Optional[Node]:
 
 
 def get_node_shape_meta(node: Node) -> Optional[torch.Tensor]:
-    return node.meta.get("val") or node.meta.get("example_value")
+    value = node.meta.get("val")
+    return value if value is not None else node.meta.get("example_value")
 
 
 def find_node_by_tag(gm: GraphModule, tag: str) -> Optional[Node]:
@@ -180,7 +181,8 @@ def find_node_by_tag(gm: GraphModule, tag: str) -> Optional[Node]:
     for node in gm.graph.nodes:
         # https://github.com/pytorch/pytorch/blob/085b71eab05cbc7d474a173884269c62d2778f77/torch/_dynamo/utils.py#L5048
         tensor_dict = node.meta.get('tensor_dict')
-        if tensor_dict and tensor_dict.get('tag') == tag:
+        node_tag = tensor_dict.get('tag') if tensor_dict else None
+        if node_tag == tag or (isinstance(node_tag, tuple) and len(node_tag) == 2 and node_tag[0] == tag):
             input_id_node = node
             break
     return input_id_node
diff --git a/deepspeed/compile/passes/sp_compile.py b/deepspeed/compile/passes/sp_compile.py
index ab2b3fb9fa33..83b30b3a9dd5 100644
--- a/deepspeed/compile/passes/sp_compile.py
+++ b/deepspeed/compile/passes/sp_compile.py
@@ -16,8 +16,9 @@
 from deepspeed.compile import constants
 
 from ..custom_ops import all_to_all, sp_dp_registry  # noqa: F401
-from ..fx import find_node_by_name, get_node_shape_meta
-from ..util import get_input_id_node, get_label_id_node, get_position_id_node, shard_tensor_node, get_sdpa_nodes
+from ..fx import get_node_shape_meta
+from ..util import (find_symbolic_shape_node, get_autosp_seq_dim, get_input_id_node, get_label_id_node,
+                    get_position_id_node, get_sdpa_nodes, shard_tensor_node)
 
 
 def prepare_autosp_inputs(input_id: torch.Tensor,
@@ -43,6 +44,8 @@ def prepare_autosp_inputs(input_id: torch.Tensor,
 
     if seq_dim < 0 or seq_dim >= input_id.ndim:
         raise ValueError(f"seq_dim {seq_dim} must be a valid index for input_id with shape {input_id.shape}")
+    if seq_dim >= label_id.ndim:
+        raise ValueError(f"seq_dim {seq_dim} is out of bounds for label_id with shape {label_id.shape}")
 
     if position_id is not None:
         if seq_dim >= position_id.ndim:
@@ -60,10 +63,10 @@ def prepare_autosp_inputs(input_id: torch.Tensor,
     if attention_mask is not None:
         torch._dynamo.decorators.mark_dynamic(attention_mask, seq_dim)
 
-    input_id.tag = constants.AUTOSP_INPUT_ID_KEY
-    label_id.tag = constants.AUTOSP_LABEL_ID_KEY
+    input_id.tag = (constants.AUTOSP_INPUT_ID_KEY, seq_dim)
+    label_id.tag = (constants.AUTOSP_LABEL_ID_KEY, seq_dim)
     if position_id is not None:
-        position_id.tag = constants.AUTOSP_POSITION_ID_KEY
+        position_id.tag = (constants.AUTOSP_POSITION_ID_KEY, seq_dim)
 
     return input_id, label_id, position_id, attention_mask
 
@@ -77,12 +80,13 @@ def pass_shard_seq_dim(gm: GraphModule, example_inputs):
 
     input_ids_node = get_input_id_node(gm)
     val = get_node_shape_meta(input_ids_node)
-    seq_symint = val.shape[1]
+    seq_dim = get_autosp_seq_dim(input_ids_node)
+    seq_symint = val.shape[seq_dim]
     assert isinstance(
         seq_symint,
         torch.SymInt), f"expected sequence dimension to be of type {torch.SymInt!r} but found {type(seq_symint)!r}"
 
-    sym_seq_dim_node = find_node_by_name(gm, str(seq_symint))
+    sym_seq_dim_node = find_symbolic_shape_node(gm, seq_symint)
     if sym_seq_dim_node is None:
         print(f"WARNING: Could not find the symbolic node for the sequence dimension")
         return
@@ -91,13 +95,10 @@ def pass_shard_seq_dim(gm: GraphModule, example_inputs):
         sharded_node = gm.graph.call_function(operator.floordiv, args=(sym_seq_dim_node, sp_size))
 
     sharded_input_nodes = set()
-    label_ids_node = get_label_id_node(gm)
     position_ids_node = get_position_id_node(gm)
 
     if input_ids_node is not None:
         sharded_input_nodes.add(input_ids_node)
-    if label_ids_node is not None:
-        sharded_input_nodes.add(label_ids_node)
     if position_ids_node is not None:
         sharded_input_nodes.add(position_ids_node)
 
@@ -133,7 +134,85 @@ def pass_shard_input_ids(gm: GraphModule, example_inputs):
 
 def pass_shard_label_ids(gm: GraphModule, example_inputs):
     label_ids_node = get_label_id_node(gm)
-    shard_tensor_node(gm, label_ids_node)
+    label_meta = get_node_shape_meta(label_ids_node)
+    seq_dim = get_autosp_seq_dim(label_ids_node)
+    seq_len = label_meta.shape[seq_dim]
+
+    def depends_on(node: Node, ancestor: Node) -> bool:
+        worklist = [node]
+        visited = set()
+        while worklist:
+            current = worklist.pop()
+            if current is ancestor:
+                return True
+            if current in visited:
+                continue
+            visited.add(current)
+            worklist.extend(current.all_input_nodes)
+        return False
+
+    loss_nodes = [node for node in gm.graph.nodes if node.target is torch.nn.functional.cross_entropy]
+    label_loss_nodes = [
+        node for node in loss_nodes
+        if len(node.args) > 1 and isinstance(node.args[1], Node) and depends_on(node.args[1], label_ids_node)
+    ]
+    if not label_loss_nodes:
+        shard_tensor_node(gm, label_ids_node, seq_dim)
+        return
+
+    causal_losses = []
+    for loss_node in label_loss_nodes:
+        weight = loss_node.kwargs.get("weight", loss_node.args[2] if len(loss_node.args) > 2 else None)
+        if weight is not None:
+            raise RuntimeError("AutoSP does not support class-weighted causal language-model cross entropy")
+        reduction = loss_node.kwargs.get("reduction", loss_node.args[6] if len(loss_node.args) > 6 else "mean")
+        if reduction != "mean":
+            raise RuntimeError(f"AutoSP only supports mean-reduced causal language-model loss, got {reduction!r}")
+        ignore_index = loss_node.kwargs.get("ignore_index", loss_node.args[4] if len(loss_node.args) > 4 else -100)
+
+        target_node = loss_node.args[1]
+        worklist = [target_node]
+        visited = set()
+        shifted_label_node = None
+        while worklist:
+            current = worklist.pop(0)
+            if current in visited or current is label_ids_node:
+                continue
+            visited.add(current)
+            current_meta = get_node_shape_meta(current)
+            if isinstance(current_meta, torch.Tensor) and current_meta.ndim == label_meta.ndim:
+                current_seq_len = current_meta.shape[seq_dim]
+                if str(current_seq_len) == str(seq_len):
+                    shifted_label_node = current
+                    break
+            worklist.extend(current.all_input_nodes)
+
+        if shifted_label_node is None:
+            shard_tensor_node(gm, label_ids_node, seq_dim)
+            return
+
+        causal_losses.append((loss_node, shifted_label_node, ignore_index))
+
+    sharded_candidates = {}
+    for loss_node, shifted_label_node, ignore_index in causal_losses:
+        if shifted_label_node not in sharded_candidates:
+            sharded_candidates[shifted_label_node] = shard_tensor_node(gm,
+                                                                       shifted_label_node,
+                                                                       seq_dim,
+                                                                       make_contiguous=True)
+        sharded_labels = sharded_candidates[shifted_label_node]
+
+        with gm.graph.inserting_after(sharded_labels):
+            valid_mask = gm.graph.call_function(operator.ne, args=(sharded_labels, ignore_index))
+        with gm.graph.inserting_after(valid_mask):
+            valid_tokens = gm.graph.call_function(torch.sum, args=(valid_mask, ))
+        with gm.graph.inserting_after(loss_node):
+            aggregated = gm.graph.call_function(torch.ops.autosp.aggregate_loss.default,
+                                                args=(loss_node, valid_tokens))
+        with gm.graph.inserting_after(aggregated):
+            global_loss = gm.graph.call_function(operator.getitem, args=(aggregated, 0))
+        loss_node.replace_all_uses_with(global_loss)
+        aggregated.update_arg(0, loss_node)
 
 
 def pass_shard_position_ids(gm: GraphModule, example_inputs):
@@ -168,6 +247,36 @@ def insert_a2a(node: Node, scatter_idx: int, gather_idx: int, name: str) -> Node
         q, k, v = attn_node.args[:3]
         suffix = f"_{idx}" if len(attention_nodes) > 1 else ""
 
+        for name, tensor in (("query", q), ("key", k), ("value", v)):
+            tensor_meta = get_node_shape_meta(tensor)
+            if tensor_meta is None or tensor_meta.ndim != 4:
+                raise RuntimeError(f"AutoSP expected a rank-4 {name} tensor for SDPA")
+            heads = tensor_meta.shape[1]
+            if isinstance(heads, int) and heads % sp_dp_registry.sp_size() != 0:
+                raise ValueError(f"AutoSP requires the {name} head count ({heads}) to be divisible by "
+                                 f"sequence_parallel_size ({sp_dp_registry.sp_size()})")
+
+        attn_mask = attn_node.kwargs.get("attn_mask")
+        mask_is_kwarg = attn_mask is not None
+        if attn_mask is None and len(attn_node.args) > 3:
+            attn_mask = attn_node.args[3]
+        if isinstance(attn_mask, Node):
+            mask_meta = get_node_shape_meta(attn_mask)
+            if mask_meta is not None and mask_meta.ndim >= 2:
+                q_meta = get_node_shape_meta(q)
+                k_meta = get_node_shape_meta(k)
+                if str(mask_meta.shape[-1]) == str(k_meta.shape[2]):
+                    raise RuntimeError("AutoSP cannot reconstruct an attention mask with a sharded key dimension. "
+                                       "Pass the full key padding mask to every SP rank.")
+                if str(mask_meta.shape[-2]) == str(q_meta.shape[2]):
+                    with gm.graph.inserting_after(attn_mask):
+                        global_mask = gm.graph.call_function(torch.ops.autosp.all_gather_sequence.default,
+                                                             args=(attn_mask, -2))
+                    if mask_is_kwarg:
+                        attn_node.update_kwarg("attn_mask", global_mask)
+                    else:
+                        attn_node.update_arg(3, global_mask)
+
         # QKV: [B, N, S/P, H] -> [B, N/P, S, H]
         insert_a2a(q, scatter_idx=1, gather_idx=2, name=f"q{suffix}")
         insert_a2a(k, scatter_idx=1, gather_idx=2, name=f"k{suffix}")
@@ -219,17 +328,27 @@ def pass_propagate_shapes(gm: torch.fx.GraphModule, real_inputs):
     saved_sdpa_masks = []
     for attn_node in get_sdpa_nodes(gm):
         attn_mask = attn_node.kwargs.get("attn_mask")
+        mask_location = "kwarg"
+        if attn_mask is None and len(attn_node.args) > 3:
+            attn_mask = attn_node.args[3]
+            mask_location = "arg"
         if attn_mask is not None:
-            saved_sdpa_masks.append((attn_node, attn_mask))
-            attn_node.update_kwarg("attn_mask", None)
+            saved_sdpa_masks.append((attn_node, mask_location, attn_mask))
+            if mask_location == "kwarg":
+                attn_node.update_kwarg("attn_mask", None)
+            else:
+                attn_node.update_arg(3, None)
 
     try:
         # fake_inputs are already created under fake_mode above, so run
         # propagation without reconverting them into a different fake mode.
         FakeTensorProp(gm, mode=fake_mode).propagate_dont_convert_inputs(*fake_inputs)
     finally:
-        for attn_node, attn_mask in saved_sdpa_masks:
-            attn_node.update_kwarg("attn_mask", attn_mask)
+        for attn_node, mask_location, attn_mask in saved_sdpa_masks:
+            if mask_location == "kwarg":
+                attn_node.update_kwarg("attn_mask", attn_mask)
+            else:
+                attn_node.update_arg(3, attn_mask)
 
 
 def apply_autosp(gm: GraphModule,
@@ -247,7 +366,7 @@ def apply_autosp(gm: GraphModule,
         debug: If True, print graph before/after each pass
         passes: Optional custom list of passes (default: DEFAULT_PASSES)
     """
-    assert sp_size * dp_size <= dist.get_world_size(), 'Insufficient device count for mesh size'
+    assert sp_size * dp_size == dist.get_world_size(), 'AutoSP mesh must cover the distributed world size'
 
     sp_dp_registry.populate_registry(sp_size, dp_size)
 
@@ -256,6 +375,7 @@ def apply_autosp(gm: GraphModule,
         pass_shard_input_ids,
         pass_shard_label_ids,
         pass_shard_position_ids,
+        pass_propagate_shapes,
         pass_insert_attention_all_to_all,
         pass_propagate_shapes,
         pass_canonicalize,
diff --git a/deepspeed/compile/util.py b/deepspeed/compile/util.py
index 5c0135ba24b8..bf3306ef02e3 100644
--- a/deepspeed/compile/util.py
+++ b/deepspeed/compile/util.py
@@ -541,6 +541,15 @@ def create_shard_offsets(gm: GraphModule, s0_node: Node) -> Tuple[Node, Node]:
     sp_size: int = sp_dp_registry.sp_size()
     sp_rank: int = dist.get_rank() % sp_dp_registry.sp_size()
     with gm.graph.inserting_after(s0_node):
+        remainder_node = gm.graph.call_function(operator.mod, args=(s0_node, sp_size))
+    with gm.graph.inserting_after(remainder_node):
+        divisible_node = gm.graph.call_function(operator.eq, args=(remainder_node, 0))
+    with gm.graph.inserting_after(divisible_node):
+        assert_node = gm.graph.call_function(
+            torch._assert,
+            args=(divisible_node, f"AutoSP sequence length must be divisible by sequence_parallel_size={sp_size}"),
+        )
+    with gm.graph.inserting_after(assert_node):
         chunk_size_node = gm.graph.call_function(operator.floordiv, args=(s0_node, sp_size))
     with gm.graph.inserting_after(chunk_size_node):
         start_node = gm.graph.call_function(operator.mul, args=(sp_rank, chunk_size_node))
@@ -593,22 +602,48 @@ def create_symbolic_slice_indices(
     return slice_all, slice_range
 
 
-def shard_tensor_node(gm: GraphModule, tensor_node: Node):
-    from .fx import find_node_by_name, get_node_shape_meta, replace_node_users
+def get_autosp_seq_dim(tensor_node: Node) -> int:
+    tensor_dict = tensor_node.meta.get("tensor_dict", {})
+    tag = tensor_dict.get("tag")
+    if isinstance(tag, tuple):
+        _, seq_dim = tag
+        return seq_dim
+    return 1
+
+
+def find_symbolic_shape_node(gm: GraphModule, symbolic_dim: torch.SymInt) -> Optional[Node]:
+    from .fx import find_node_by_name
+    node = find_node_by_name(gm, str(symbolic_dim))
+    if node is not None:
+        return node
+
+    for candidate in gm.graph.nodes:
+        candidate_value = candidate.meta.get("val", candidate.meta.get("example_value"))
+        if candidate.op == "placeholder" and str(candidate_value) == str(symbolic_dim):
+            return candidate
+    return None
+
+
+def shard_tensor_node(gm: GraphModule,
+                      tensor_node: Node,
+                      seq_dim: Optional[int] = None,
+                      make_contiguous: bool = False) -> Node:
+    from .fx import get_node_shape_meta, replace_node_users
     val = get_node_shape_meta(tensor_node)
     assert val is not None, f"Node {tensor_node.name} has no shape metadata"
 
-    seq_len = val.shape[1]
+    seq_dim = get_autosp_seq_dim(tensor_node) if seq_dim is None else seq_dim
+    seq_len = val.shape[seq_dim]
 
     assert isinstance(
         seq_len,
         torch.SymInt), (f"Expected sequence dimension to be {torch.SymInt!r} but instead found {type(seq_len)!r}")
 
-    symb_seq_int_node = find_node_by_name(gm, str(seq_len))
+    symb_seq_int_node = find_symbolic_shape_node(gm, seq_len)
     assert symb_seq_int_node, f"Unable to find symbolic placeholder for {seq_len}"
 
     slice_all, slice_range = create_symbolic_slice_indices(gm, symb_seq_int_node)
-    indices = (slice_all, slice_range)
+    indices = tuple(slice_range if dim == seq_dim else slice_all for dim in range(val.ndim))
 
     positions = {node: i for i, node in enumerate(gm.graph.nodes)}
     # Insert after the later dependency so the new getitem does not appear
@@ -622,3 +657,9 @@ def shard_tensor_node(gm: GraphModule, tensor_node: Node):
         )
 
     replace_node_users(tensor_node, sliced_node, exclude=[sliced_node])
+    if make_contiguous:
+        with gm.graph.inserting_after(sliced_node):
+            contiguous_node = gm.graph.call_method("contiguous", args=(sliced_node, ))
+        replace_node_users(sliced_node, contiguous_node, exclude=[contiguous_node])
+        return contiguous_node
+    return sliced_node
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 86918bd71c5a..0a9ac6161af5 100644
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -945,7 +945,7 @@ def destroy(self):
         optimizer = getattr(self, "optimizer", None)
         if optimizer is not None and hasattr(optimizer, 'destroy'):
             optimizer.destroy()
-        if self.is_deepcompile_active():
+        if self.is_deepcompile_active() and not self.uses_parallelization_pass_only():
             get_deepcompile_handle().cleanup()
         debug_clear_module_and_param_names()
 
@@ -3027,7 +3027,7 @@ def _backward_prologue(self):
         assert not self.eigenvalue_enabled(), "Eigenvalue is not supported with non-scalar backward"
         assert not self.amp_enabled(), "Apex AMP is not supported with non-scalar backward"
 
-        if self.is_deepcompile_active() and not self.compile_autotp():
+        if self.is_deepcompile_active() and not self.uses_parallelization_pass_only():
             deepcompile_backward_prologue(self.is_gradient_accumulation_boundary())
 
         if isinstance(self.optimizer, ZeROOptimizer):
@@ -3062,7 +3062,7 @@ def _backward_epilogue(self):
                 self.optimizer.backward_epilogue()
             self.optimizer.exit_backward()
 
-        if self.is_deepcompile_active() and not self.compile_autotp():
+        if self.is_deepcompile_active() and not self.uses_parallelization_pass_only():
             deepcompile_backward_epilogue()
 
         see_memory_usage("Engine after backward", force=self.memory_breakdown())
diff --git a/deepspeed/sequence/fpdt_layer.py b/deepspeed/sequence/fpdt_layer.py
index f366ce40425e..b680be1f8a1f 100644
--- a/deepspeed/sequence/fpdt_layer.py
+++ b/deepspeed/sequence/fpdt_layer.py
@@ -358,8 +358,8 @@ def backward(ctx, grad_output):
         del grad_output
 
         dq = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)]
-        dk = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)]
-        dv = [torch.zeros(global_q[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)]
+        dk = [torch.zeros(global_k[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)]
+        dv = [torch.zeros(global_v[0].shape, dtype=torch.float, device=device) for _ in range(num_chunks)]
 
         grad_qkv_linear_weight = torch.zeros(qkv_linear_weight.shape,
                                              device=qkv_linear_weight.device,
diff --git a/docs/code-docs/source/training.rst b/docs/code-docs/source/training.rst
index 63c824f88f72..a36884a45c24 100644
--- a/docs/code-docs/source/training.rst
+++ b/docs/code-docs/source/training.rst
@@ -719,6 +719,7 @@ config and calling ``prepare_autosp_inputs()`` to prepare inputs before each for
 .. note::
    AutoSP requires ZeRO stage 0 (no ZeRO optimization). Using AutoSP with ZeRO stages 1, 2, or 3 is not currently supported.
    AutoSP also requires ``torch.nn.functional.scaled_dot_product_attention()`` as the attention backend.
+   The sequence length and every SDPA Q/K/V head count must be divisible by ``sequence_parallel_size``.
 
 Input Preparation
 ~~~~~~~~~~~~~~~~~
@@ -740,6 +741,7 @@ automatic sharding:
     )
 
 This serves as a hint to the compiler to know which inputs should be sharded across which dimension.
+``seq_dim`` is preserved in the compiler metadata, so both batch-first and sequence-first input layouts are supported.
 
 Memory Optimization
 ~~~~~~~~~~~~~~~~~~~
diff --git a/tests/unit/v1/compile/test_compile_autosp.py b/tests/unit/v1/compile/test_compile_autosp.py
index bcce3ed9f798..579424912b76 100644
--- a/tests/unit/v1/compile/test_compile_autosp.py
+++ b/tests/unit/v1/compile/test_compile_autosp.py
@@ -26,18 +26,80 @@
 _SP_SIZE = 2
 
 
-def _create_sdpa_graph(seq_len):
+def _create_sdpa_graph(seq_len, num_heads=2, mask_rank=None):
     graph = Graph()
     inputs = []
     for name in ("query", "key", "value"):
         node = graph.placeholder(name)
-        node.meta["example_value"] = torch.empty(1, 2, seq_len, 8)
+        node.meta["example_value"] = torch.empty(1, num_heads, seq_len, 8)
         inputs.append(node)
-    sdpa = graph.call_function(F.scaled_dot_product_attention, args=tuple(inputs))
+    kwargs = {}
+    if mask_rank is not None:
+        mask = graph.placeholder("attention_mask")
+        mask_shape = (seq_len, seq_len * _SP_SIZE)
+        if mask_rank == 4:
+            mask_shape = (1, 1) + mask_shape
+        mask.meta["example_value"] = torch.empty(mask_shape)
+        kwargs["attn_mask"] = mask
+    sdpa = graph.call_function(F.scaled_dot_product_attention, args=tuple(inputs), kwargs=kwargs)
     graph.output(sdpa)
     return GraphModule({}, graph)
 
 
+def _create_causal_loss_graph(seq_len=16, ignore_index=-100, shift_labels=True):
+
+    class CausalLoss(torch.nn.Module):
+
+        def forward(self, input_ids, labels):
+            logits = F.one_hot(input_ids, num_classes=32).float()
+            targets = F.pad(labels, (0, 1), value=ignore_index)[..., 1:].contiguous() if shift_labels else labels
+            return F.cross_entropy(logits.view(-1, 32), targets.view(-1), ignore_index=ignore_index)
+
+    torch._dynamo.reset()
+    input_ids = torch.randint(0, 32, (2, seq_len))
+    labels = input_ids.clone()
+    input_ids.tag = constants.AUTOSP_INPUT_ID_KEY
+    labels.tag = constants.AUTOSP_LABEL_ID_KEY
+    torch._dynamo.decorators.mark_dynamic(input_ids, 1)
+    torch._dynamo.decorators.mark_dynamic(labels, 1)
+
+    captured_gm = [None]
+
+    def capture(gm, example_inputs):
+        captured_gm[0] = gm
+        return gm
+
+    compiled = torch.compile(CausalLoss(), backend=capture, dynamic=True)
+    compiled(input_ids, labels)
+    return captured_gm[0]
+
+
+def _create_sequence_first_graph(seq_len=16):
+
+    class SequenceFirst(torch.nn.Module):
+
+        def forward(self, input_ids, labels):
+            return input_ids.float().sum() + labels.float().sum()
+
+    torch._dynamo.reset()
+    input_ids = torch.ones(seq_len, 2, dtype=torch.long)
+    labels = input_ids.clone()
+    input_ids.tag = (constants.AUTOSP_INPUT_ID_KEY, 0)
+    labels.tag = (constants.AUTOSP_LABEL_ID_KEY, 0)
+    torch._dynamo.decorators.mark_dynamic(input_ids, 0)
+    torch._dynamo.decorators.mark_dynamic(labels, 0)
+
+    captured_gm = [None]
+
+    def capture(gm, example_inputs):
+        captured_gm[0] = gm
+        return gm
+
+    compiled = torch.compile(SequenceFirst(), backend=capture, dynamic=True)
+    compiled(input_ids, labels)
+    return captured_gm[0]
+
+
 class TestAutoSPCompile(DistributedTest):
     world_size = 4
     non_daemonic_procs = True
@@ -171,6 +233,14 @@ def test(self, seq_len):
         # create_shard_offsets emits: chunk = seq // sp_size; start = rank * chunk; end = start + chunk.
         # Verify the three-node chain has the right operators and wiring.
         chunk_size_node = start_node.args[1]  # start = rank * chunk  →  chunk is arg[1]
+        assert_node = next(node for node in gm.graph.nodes if node.target == torch._assert)
+        divisible_node = assert_node.args[0]
+        remainder_node = divisible_node.args[0]
+
+        assert divisible_node.target == operator.eq
+        assert divisible_node.args[1] == 0
+        assert remainder_node.target == operator.mod
+        assert remainder_node.args == (sym_seq_node, _SP_SIZE)
 
         assert chunk_size_node.target == operator.floordiv
         assert chunk_size_node.args[0] is sym_seq_node
@@ -304,3 +374,188 @@ def test_preserves_topological_order_when_sym_placeholder_follows_input(self):
             shard_tensor_node(reordered_gm, reordered_input_ids)
 
         reordered_gm.graph.lint()
+
+
+class TestAutoSPValidation:
+
+    def test_reads_tensor_shape_metadata_without_boolean_conversion(self):
+        from deepspeed.compile.fx import get_node_shape_meta
+
+        node = Graph().placeholder("input")
+        value = torch.empty(2)
+        node.meta["val"] = value
+
+        assert get_node_shape_meta(node) is value
+
+    def test_prepare_inputs_preserves_non_default_sequence_dimension(self):
+        from deepspeed.compile.passes.sp_compile import prepare_autosp_inputs
+
+        input_ids = torch.ones(8, 2, dtype=torch.long)
+        labels = input_ids.clone()
+        position_ids = input_ids.clone()
+        with patch.object(torch._dynamo.decorators, "mark_dynamic"):
+            prepare_autosp_inputs(input_ids, labels, position_ids, seq_dim=0)
+
+        assert input_ids.tag == (constants.AUTOSP_INPUT_ID_KEY, 0)
+        assert labels.tag == (constants.AUTOSP_LABEL_ID_KEY, 0)
+        assert position_ids.tag == (constants.AUTOSP_POSITION_ID_KEY, 0)
+
+    @pytest.mark.sequential
+    def test_shards_non_default_sequence_dimension(self):
+        import deepspeed.comm as _dist
+        from deepspeed.compile.custom_ops import sp_dp_registry
+        from deepspeed.compile.passes.sp_compile import pass_shard_input_ids
+        from deepspeed.compile.util import get_input_id_node
+
+        gm = _create_sequence_first_graph()
+        input_node = get_input_id_node(gm)
+        with patch.object(sp_dp_registry, "sp_size", return_value=_SP_SIZE), \
+             patch.object(_dist, "get_rank", return_value=0):
+            pass_shard_input_ids(gm, ())
+
+        shard = next(node for node in gm.graph.nodes if node.target == operator.getitem and node.args[0] is input_node)
+        indices = shard.args[1]
+        assert indices[0].target == slice
+        assert indices[0].args[0] is not None
+        assert indices[1].target == slice
+        assert indices[1].args == (None, None, None)
+
+    @pytest.mark.sequential
+    def test_rejects_non_divisible_sequence_length(self):
+        import deepspeed.comm as _dist
+        from deepspeed.compile.custom_ops import sp_dp_registry
+        from deepspeed.compile.passes.sp_compile import pass_canonicalize, pass_shard_input_ids
+
+        gm = _create_sequence_first_graph()
+        with patch.object(sp_dp_registry, "sp_size", return_value=_SP_SIZE), \
+             patch.object(_dist, "get_rank", return_value=0):
+            pass_shard_input_ids(gm, ())
+            pass_canonicalize(gm, ())
+
+        input_ids = torch.ones(15, 2, dtype=torch.long)
+        labels = input_ids.clone()
+        with pytest.raises(AssertionError, match="sequence length must be divisible"):
+            gm(15, 2, input_ids, 15, labels)
+
+    def test_rejects_changed_mesh(self):
+        from deepspeed.compile.custom_ops import sp_dp_registry
+
+        registry = {"SP_SIZE": 2, "DP_SIZE": 2, "is_reg": True}
+        with patch.object(sp_dp_registry, "GROUP_REGISTRY", registry), \
+             patch.object(sp_dp_registry.dist, "get_world_size", return_value=4):
+            with pytest.raises(RuntimeError, match="already initialized"):
+                sp_dp_registry.populate_registry(4, 1)
+
+    def test_rejects_partial_mesh(self):
+        from deepspeed.compile.custom_ops import sp_dp_registry
+
+        with patch.object(sp_dp_registry.dist, "get_world_size", return_value=4):
+            with pytest.raises(ValueError, match="must cover"):
+                sp_dp_registry.populate_registry(2, 1)
+
+    def test_all_ignored_loss_shard_contributes_zero(self):
+        import deepspeed.comm as _dist
+        from deepspeed.compile.custom_ops import sp_dp_registry
+        from deepspeed.compile.custom_ops.all_to_all import aggregate_loss
+
+        loss = torch.tensor(float("nan"), requires_grad=True)
+        valid_tokens = torch.tensor(0)
+        registry = {0: object(), "SP_SIZE": 2, "DP_SIZE": 1, "is_reg": True}
+        with patch.object(sp_dp_registry, "GROUP_REGISTRY", registry), \
+             patch.object(_dist, "get_rank", return_value=0), \
+             patch.object(_dist, "all_reduce"):
+            global_loss, weight = aggregate_loss(loss, valid_tokens)
+            global_loss.backward()
+
+        assert global_loss.item() == 0
+        assert weight.item() == 0
+        assert loss.grad.item() == 0
+
+    def test_rejects_non_divisible_attention_heads(self):
+        from deepspeed.compile.custom_ops import sp_dp_registry
+        from deepspeed.compile.passes.sp_compile import pass_insert_attention_all_to_all
+
+        gm = _create_sdpa_graph(seq_len=8, num_heads=3)
+        with patch.object(sp_dp_registry, "sp_size", return_value=_SP_SIZE):
+            with pytest.raises(ValueError, match="query head count"):
+                pass_insert_attention_all_to_all(gm, ())
+
+    def test_gathers_local_attention_mask_query_dimension(self):
+        from deepspeed.compile.custom_ops import sp_dp_registry
+        from deepspeed.compile.passes.sp_compile import pass_insert_attention_all_to_all
+
+        gm = _create_sdpa_graph(seq_len=8, mask_rank=4)
+        with patch.object(sp_dp_registry, "sp_size", return_value=_SP_SIZE):
+            pass_insert_attention_all_to_all(gm, ())
+
+        sdpa_node = next(node for node in gm.graph.nodes if node.target == F.scaled_dot_product_attention)
+        mask_node = sdpa_node.kwargs["attn_mask"]
+        assert mask_node.target == torch.ops.autosp.all_gather_sequence.default
+        assert mask_node.args[1] == -2
+
+    def test_gathers_rank_two_attention_mask_query_dimension(self):
+        from deepspeed.compile.custom_ops import sp_dp_registry
+        from deepspeed.compile.passes.sp_compile import pass_insert_attention_all_to_all
+
+        gm = _create_sdpa_graph(seq_len=8, mask_rank=2)
+        with patch.object(sp_dp_registry, "sp_size", return_value=_SP_SIZE):
+            pass_insert_attention_all_to_all(gm, ())
+
+        sdpa_node = next(node for node in gm.graph.nodes if node.target == F.scaled_dot_product_attention)
+        assert sdpa_node.kwargs["attn_mask"].target == torch.ops.autosp.all_gather_sequence.default
+
+    @pytest.mark.sequential
+    def test_shards_labels_after_causal_shift(self):
+        import deepspeed.comm as _dist
+        from deepspeed.compile.custom_ops import sp_dp_registry
+        from deepspeed.compile.passes.sp_compile import pass_shard_label_ids
+        from deepspeed.compile.util import get_label_id_node
+
+        gm = _create_causal_loss_graph(seq_len=64)
+        label_node = get_label_id_node(gm)
+        with patch.object(sp_dp_registry, "sp_size", return_value=_SP_SIZE), \
+             patch.object(_dist, "get_rank", return_value=0):
+            pass_shard_label_ids(gm, ())
+
+        raw_label_slices = [
+            node for node in gm.graph.nodes if node.target == operator.getitem and node.args[0] is label_node
+        ]
+        assert not raw_label_slices
+        assert any(node.target == torch.ops.autosp.aggregate_loss.default for node in gm.graph.nodes)
+        gm.graph.lint()
+
+    @pytest.mark.sequential
+    def test_shards_direct_cross_entropy_labels(self):
+        import deepspeed.comm as _dist
+        from deepspeed.compile.custom_ops import sp_dp_registry
+        from deepspeed.compile.passes.sp_compile import pass_shard_label_ids
+        from deepspeed.compile.util import get_label_id_node
+
+        gm = _create_causal_loss_graph(seq_len=64, shift_labels=False)
+        label_node = get_label_id_node(gm)
+        with patch.object(sp_dp_registry, "sp_size", return_value=_SP_SIZE), \
+             patch.object(_dist, "get_rank", return_value=0):
+            pass_shard_label_ids(gm, ())
+
+        raw_label_slices = [
+            node for node in gm.graph.nodes if node.target == operator.getitem and node.args[0] is label_node
+        ]
+        assert len(raw_label_slices) == 1
+        assert not any(node.target == torch.ops.autosp.aggregate_loss.default for node in gm.graph.nodes)
+        gm.graph.lint()
+
+    @pytest.mark.sequential
+    def test_causal_loss_uses_configured_ignore_index(self):
+        import deepspeed.comm as _dist
+        from deepspeed.compile.custom_ops import sp_dp_registry
+        from deepspeed.compile.passes.sp_compile import pass_shard_label_ids
+
+        gm = _create_causal_loss_graph(seq_len=64, ignore_index=-1)
+        with patch.object(sp_dp_registry, "sp_size", return_value=_SP_SIZE), \
+             patch.object(_dist, "get_rank", return_value=0):
+            pass_shard_label_ids(gm, ())
+
+        loss_node = next(node for node in gm.graph.nodes if node.target is F.cross_entropy)
+        ignore_index = loss_node.kwargs["ignore_index"] if "ignore_index" in loss_node.kwargs else loss_node.args[4]
+        valid_mask = next(node for node in gm.graph.nodes if node.target == operator.ne)
+        assert valid_mask.args[1] is ignore_index
diff --git a/tests/unit/v1/compile/util.py b/tests/unit/v1/compile/util.py
index c61554091a1c..ce6caa0a0b47 100644
--- a/tests/unit/v1/compile/util.py
+++ b/tests/unit/v1/compile/util.py
@@ -136,9 +136,12 @@ def compare_sp_loss(self, config, sp_size, iterations=3):
     # AutoSP's graph pass can therefore find F.scaled_dot_product_attention nodes.
     def _sdpa_inner(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, scale=None):
         # DistributedAttention delivers tensors in [b, s, n, h]; SDPA wants [b, n, s, h].
+        if attn_mask is not None and attn_mask.ndim >= 3:
+            attn_mask = torch.ops.autosp.all_gather_sequence.default(attn_mask, -2)
         out = F.scaled_dot_product_attention(q.permute(0, 2, 1, 3),
                                              k.permute(0, 2, 1, 3),
                                              v.permute(0, 2, 1, 3),
+                                             attn_mask=attn_mask,
                                              dropout_p=dropout_p,
                                              is_causal=is_causal,
                                              scale=scale)
@@ -158,7 +161,14 @@ def _ulysses_attn_forward(module,
         q = query_states.transpose(1, 2).contiguous()
         k = key_states.transpose(1, 2).contiguous()
         v = value_states.transpose(1, 2).contiguous()
-        out = _dist_attn(q, k, v, batch_dim_idx=0, dropout_p=dropout, is_causal=is_causal, scale=scaling)
+        out = _dist_attn(q,
+                         k,
+                         v,
+                         batch_dim_idx=0,
+                         attn_mask=attention_mask,
+                         dropout_p=dropout,
+                         is_causal=is_causal,
+                         scale=scaling)
         return out, None
 
     ALL_ATTENTION_FUNCTIONS["ulyssess"] = _ulysses_attn_forward
@@ -187,26 +197,35 @@ def _ulysses_attn_forward(module,
     for i in range(iterations):
         torch.manual_seed(42 + i)
         full_ids = torch.randint(0, vocab_size, (1, seq_length), device=device)
+        full_mask = torch.ones_like(full_ids)
+        full_mask[:, -4:] = 0
+        full_labels = full_ids.masked_fill(full_mask == 0, -100)
+        shifted_labels = F.pad(full_labels, (0, 1), value=-100)[..., 1:]
 
         # Ulysses: each rank processes its own shard.
         shard_ids = full_ids[:, sp_rank * chunk:(sp_rank + 1) * chunk]
+        shard_labels = full_labels[:, sp_rank * chunk:(sp_rank + 1) * chunk]
+        shard_shifted_labels = shifted_labels[:, sp_rank * chunk:(sp_rank + 1) * chunk]
         shard_pos = torch.arange(sp_rank * chunk, (sp_rank + 1) * chunk, device=device).unsqueeze(0)
-        shard_mask = torch.ones(1, chunk, device=device, dtype=torch.long)
+        shard_mask = full_mask
         ul_out = ulysses_engine(input_ids=shard_ids,
-                                labels=shard_ids,
+                                labels=shard_labels,
+                                shift_labels=shard_shifted_labels,
                                 position_ids=shard_pos,
                                 attention_mask=shard_mask)
-        # Average per-shard losses across SP ranks to get the full-sequence loss.
-        ul_loss = ul_out.loss.clone()
+        local_valid_tokens = (shard_shifted_labels != -100).sum()
+        total_valid_tokens = local_valid_tokens.clone()
+        dist.all_reduce(total_valid_tokens, group=sp_group)
+        weighted_ul_loss = ul_out.loss * local_valid_tokens / total_valid_tokens.clamp_min(1)
+        ul_loss = weighted_ul_loss.detach().clone()
         dist.all_reduce(ul_loss, group=sp_group)
-        ul_loss = ul_loss / sp_size
 
         # AutoSP: full sequence.  dynamic=True makes all shapes symbolic, so mark_dynamic
         # is not needed; only the tag attributes that the autosp pass uses are set here.
         autosp_ids = full_ids.clone()
-        autosp_lbl = autosp_ids.clone()
+        autosp_lbl = full_labels.clone()
         autosp_pos = torch.arange(seq_length, device=device).unsqueeze(0)
-        autosp_msk = torch.ones(1, seq_length, device=device, dtype=torch.long)
+        autosp_msk = full_mask.clone()
         autosp_ids.tag = autosp_constants.AUTOSP_INPUT_ID_KEY
         autosp_lbl.tag = autosp_constants.AUTOSP_LABEL_ID_KEY
         autosp_pos.tag = autosp_constants.AUTOSP_POSITION_ID_KEY
@@ -216,7 +235,7 @@ def _ulysses_attn_forward(module,
                                    attention_mask=autosp_msk)
         autosp_loss = autosp_out.loss
 
-        ulysses_engine.backward(ul_out.loss)
+        ulysses_engine.backward(weighted_ul_loss)
         ulysses_engine.step()
         autosp_engine.backward(autosp_loss)
         autosp_engine.step()