diff --git a/.github/workflows/nv-flash-attn.yml b/.github/workflows/nv-flash-attn.yml
index 69f6988aa415..61d6768b866a 100644
--- a/.github/workflows/nv-flash-attn.yml
+++ b/.github/workflows/nv-flash-attn.yml
@@ -5,7 +5,7 @@ on:
   pull_request:
     paths:
       - 'deepspeed/sequence/**'
-      - 'tests/unit/sequence_parallelism/**'
+      - 'tests/unit/v1/sequence_parallelism/**'
       - '.github/workflows/nv-flash-attn.yml'
   schedule:
     - cron: "0 0 * * *"
@@ -57,7 +57,7 @@ jobs:
         run: |
           unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
           cd tests
-          python -m pytest --color=yes --durations=0 --verbose -rF unit/sequence_parallelism/test_ulysses.py --torch_ver="2.6" --cuda_ver="12"
+          python -m pytest --color=yes --durations=0 --verbose -rF unit/v1/sequence_parallelism/test_ulysses.py --torch_ver="2.6" --cuda_ver="12"
       - name: Open GitHub issue if nightly CI fails
         if: ${{ failure() && (github.event_name == 'schedule') }}
         uses: JasonEtco/create-an-issue@v2
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 1ade9d9ee37a..86918bd71c5a 100644
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1923,7 +1923,7 @@ def _configure_distributed_model(self, model):
                     "DeepSpeed Sequence Parallelism (Ulysses) with PyTorch < 2.3 may encounter "
                     "rank indexing errors during backward pass when sp_size < world_size. "
                     "Please use the weighted all-reduce workaround shown in the regression test "
-                    "(https://github.com/deepspeedai/DeepSpeed/blob/master/tests/unit/sequence_parallelism/test_ulysses.py) "
+                    "(https://github.com/deepspeedai/DeepSpeed/blob/master/tests/unit/v1/sequence_parallelism/test_ulysses.py) "
                     "or upgrade to PyTorch 2.3+.")
             self.communication_data_type = self._config.seq_parallel_communication_data_type
             self.seq_parallel_group = groups._get_sequence_parallel_group()
diff --git a/deepspeed/sequence/fpdt_layer.py b/deepspeed/sequence/fpdt_layer.py
index 938a7e419a0b..f366ce40425e 100644
--- a/deepspeed/sequence/fpdt_layer.py
+++ b/deepspeed/sequence/fpdt_layer.py
@@ -155,6 +155,9 @@ def forward(ctx: Any,
                 cpu_offloading=True):
 
         do_save = layernorm_output.requires_grad
+        # The all-to-all partitions by KV group, and only the caller knows that count: the
+        # gather-direction and backward calls below see an already-sharded head dim.
+        num_kv_heads = kv_projection_size // hidden_size_per_attention_head
 
         if rotary_pos_emb is not None:
             pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3)
@@ -180,6 +183,7 @@ def forward(ctx: Any,
             ctx.dtype = layernorm_output.dtype
             ctx.projection_size = projection_size
             ctx.kv_projection_size = kv_projection_size
+            ctx.num_kv_heads = num_kv_heads
 
             global_q = []
             global_k = []
@@ -206,7 +210,7 @@ def forward(ctx: Any,
                 q_chunk = qkv_chunk[:, :, :projection_size].contiguous().reshape(
                     qkv_chunk.shape[0], qkv_chunk.shape[1], -1,
                     hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous()  # b, l, nh, hd
-                q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg)
+                q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg, num_kv_heads=num_kv_heads)
                 global_q_chunk_len = q_chunk.shape[1]
                 if rotary_pos_emb is not None:
                     q_chunk = apply_rotary_pos_emb(q_chunk,
@@ -217,7 +221,7 @@ def forward(ctx: Any,
                 k_chunk = qkv_chunk[:, :, projection_size:projection_size + kv_projection_size].contiguous().reshape(
                     qkv_chunk.shape[0], qkv_chunk.shape[1], -1,
                     hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous()  # b, l, nh, hd
-                k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg)
+                k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg, num_kv_heads=num_kv_heads)
                 if rotary_pos_emb is not None:
                     k_chunk = apply_rotary_pos_emb(k_chunk,
                                                    pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)],
@@ -227,7 +231,7 @@ def forward(ctx: Any,
                 v_chunk = qkv_chunk[:, :, projection_size + kv_projection_size:].contiguous().reshape(
                     qkv_chunk.shape[0], qkv_chunk.shape[1], -1,
                     hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous()  # b, l, nh, hd
-                v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg)
+                v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg, num_kv_heads=num_kv_heads)
                 global_v.append(v_chunk)
 
                 for k_i in range(len(global_k)):
@@ -276,7 +280,12 @@ def forward(ctx: Any,
 
             for i in range(num_chunks):
                 global_lse[i] = global_lse[i][:, :, :, 0].permute(0, 2, 1).contiguous()
-                output[i] = single_all_to_all(global_o[i].to(ctx.dtype).contiguous(), gather_idx, scatter_idx, 0, spg)
+                output[i] = single_all_to_all(global_o[i].to(ctx.dtype).contiguous(),
+                                              gather_idx,
+                                              scatter_idx,
+                                              0,
+                                              spg,
+                                              num_kv_heads=num_kv_heads)
             output = torch.cat(output, dim=1)
 
             head_dim = output.shape[-1]
@@ -312,6 +321,7 @@ def backward(ctx, grad_output):
 
         projection_size = ctx.projection_size
         kv_projection_size = ctx.kv_projection_size
+        num_kv_heads = ctx.num_kv_heads
 
         layernorm_output = ctx.saved_tensors[0]
 
@@ -338,7 +348,12 @@ def backward(ctx, grad_output):
             st = chunk_size * i
             ed = st + chunk_size
             grad_global_attn_output.append(
-                single_all_to_all(grad_output[:, st:ed].contiguous(), scatter_idx, gather_idx, 0, spg))
+                single_all_to_all(grad_output[:, st:ed].contiguous(),
+                                  scatter_idx,
+                                  gather_idx,
+                                  0,
+                                  spg,
+                                  num_kv_heads=num_kv_heads))
 
         del grad_output
 
@@ -440,8 +455,8 @@ def backward(ctx, grad_output):
             else:
                 dk[i] = dk[i].to(dtype)
             dv[i] = dv[i].to(dtype)
-            dk[i] = single_all_to_all(dk[i].contiguous(), gather_idx, scatter_idx, 0, spg)
-            dv[i] = single_all_to_all(dv[i].contiguous(), gather_idx, scatter_idx, 0, spg)
+            dk[i] = single_all_to_all(dk[i].contiguous(), gather_idx, scatter_idx, 0, spg, num_kv_heads=num_kv_heads)
+            dv[i] = single_all_to_all(dv[i].contiguous(), gather_idx, scatter_idx, 0, spg, num_kv_heads=num_kv_heads)
 
             input_st = i * input_chunk_size
             input_ed = input_st + input_chunk_size
@@ -474,7 +489,12 @@ def backward(ctx, grad_output):
                                                       ctx.pos_emb_sin[:, dq_seq_len * i:dq_seq_len * (i + 1)])
             else:
                 dq[i] = dq[i].to(dtype)
-            dq[i] = single_all_to_all(dq[i].to(dtype).contiguous(), gather_idx, scatter_idx, 0, spg)
+            dq[i] = single_all_to_all(dq[i].to(dtype).contiguous(),
+                                      gather_idx,
+                                      scatter_idx,
+                                      0,
+                                      spg,
+                                      num_kv_heads=num_kv_heads)
 
             input_chunk = layernorm_output[:input_chunk_size].reshape(-1, layernorm_output.shape[-1])
             layernorm_output = layernorm_output[input_chunk_size:]
@@ -565,6 +585,9 @@ def forward(ctx: Any,
                 cpu_offloading=True):
 
         do_save = layernorm_output.requires_grad
+        # The all-to-all partitions by KV group, and only the caller knows that count: the
+        # gather-direction and backward calls below see an already-sharded head dim.
+        num_kv_heads = kv_projection_size // hidden_size_per_attention_head
 
         if rotary_pos_emb is not None:
             pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3)
@@ -590,6 +613,7 @@ def forward(ctx: Any,
             ctx.dtype = layernorm_output.dtype
             ctx.projection_size = projection_size
             ctx.kv_projection_size = kv_projection_size
+            ctx.num_kv_heads = num_kv_heads
 
             global_q = []
             global_k = []
@@ -628,18 +652,18 @@ def forward(ctx: Any,
                 q_chunk = qkv_chunk[:, :, :projection_size].contiguous().reshape(
                     qkv_chunk.shape[0], qkv_chunk.shape[1], -1,
                     hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous()  # b, l, nh, hd
-                q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg)
+                q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg, num_kv_heads=num_kv_heads)
                 global_q_chunk_len = q_chunk.shape[1]
 
                 k_chunk = qkv_chunk[:, :, projection_size:projection_size + kv_projection_size].contiguous().reshape(
                     qkv_chunk.shape[0], qkv_chunk.shape[1], -1,
                     hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous()  # b, l, nh, hd
-                k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg)
+                k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg, num_kv_heads=num_kv_heads)
 
                 v_chunk = qkv_chunk[:, :, projection_size + kv_projection_size:].contiguous().reshape(
                     qkv_chunk.shape[0], qkv_chunk.shape[1], -1,
                     hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous()  # b, l, nh, hd
-                v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg)
+                v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg, num_kv_heads=num_kv_heads)
 
                 dist.barrier()
 
@@ -739,8 +763,12 @@ def forward(ctx: Any,
                 global_q[q_compute_chunk_idx].offload()
                 q_compute_chunk_idx += 1
 
-                all2all_output = single_all_to_all(
-                    cur_attn_output.to(ctx.dtype).contiguous(), gather_idx, scatter_idx, 0, spg)
+                all2all_output = single_all_to_all(cur_attn_output.to(ctx.dtype).contiguous(),
+                                                   gather_idx,
+                                                   scatter_idx,
+                                                   0,
+                                                   spg,
+                                                   num_kv_heads=num_kv_heads)
                 final_output.append(all2all_output)
                 with get_accelerator().stream(general_offload_stream):
                     global_o.append(SequenceChunk(cur_attn_output.to(ctx.dtype)))
@@ -783,6 +811,7 @@ def backward(ctx, grad_output):
 
         projection_size = ctx.projection_size
         kv_projection_size = ctx.kv_projection_size
+        num_kv_heads = ctx.num_kv_heads
 
         layernorm_output = ctx.layernorm_output
 
@@ -821,8 +850,12 @@ def backward(ctx, grad_output):
                                                device=qkv_linear_weight.device,
                                                dtype=torch.float)
 
-        grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(), scatter_idx,
-                                                          gather_idx, 0, spg)
+        grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(),
+                                                          scatter_idx,
+                                                          gather_idx,
+                                                          0,
+                                                          spg,
+                                                          num_kv_heads=num_kv_heads)
         get_accelerator().synchronize()
         grad_output = grad_output[:, chunk_size:]
 
@@ -931,7 +964,11 @@ def backward(ctx, grad_output):
 
                         if grad_global_attn_output[next_q_compute_chunk_idx] is None:
                             grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(),
-                                                                              scatter_idx, gather_idx, 0, spg)
+                                                                              scatter_idx,
+                                                                              gather_idx,
+                                                                              0,
+                                                                              spg,
+                                                                              num_kv_heads=num_kv_heads)
                             dist.barrier()
                             grad_output = grad_output[:, chunk_size:]
                             grad_global_attn_output[next_q_compute_chunk_idx] = SequenceChunk(
@@ -976,9 +1013,24 @@ def backward(ctx, grad_output):
                 dk_accum = dk_accum.to(dtype)
             dv_accum = dv_accum.to(dtype)
 
-            dq_accum = single_all_to_all(dq_accum.contiguous(), gather_idx, scatter_idx, 0, spg)
-            dk_accum = single_all_to_all(dk_accum.contiguous(), gather_idx, scatter_idx, 0, spg)
-            dv_accum = single_all_to_all(dv_accum.contiguous(), gather_idx, scatter_idx, 0, spg)
+            dq_accum = single_all_to_all(dq_accum.contiguous(),
+                                         gather_idx,
+                                         scatter_idx,
+                                         0,
+                                         spg,
+                                         num_kv_heads=num_kv_heads)
+            dk_accum = single_all_to_all(dk_accum.contiguous(),
+                                         gather_idx,
+                                         scatter_idx,
+                                         0,
+                                         spg,
+                                         num_kv_heads=num_kv_heads)
+            dv_accum = single_all_to_all(dv_accum.contiguous(),
+                                         gather_idx,
+                                         scatter_idx,
+                                         0,
+                                         spg,
+                                         num_kv_heads=num_kv_heads)
 
             general_offload_stream.synchronize()
             compute_stream.wait_stream(general_offload_stream)
diff --git a/deepspeed/sequence/layer.py b/deepspeed/sequence/layer.py
index 5d3b1d95f488..0f828b39148d 100644
--- a/deepspeed/sequence/layer.py
+++ b/deepspeed/sequence/layer.py
@@ -15,27 +15,28 @@
 from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list
 from deepspeed.utils import groups
 
-# Ulysses sequence parallelism keeps its own kv-head count, memoized on the first uneven
-# all-to-all. The state lives here, independent of AutoTP.
-# TODO: this is process-wide and never reset, so the first model to take the uneven path locks
-# every later Ulysses call into it -- a second model with a different head count would then be
-# split against the first one's value. Moving it onto the attention instance means threading the
-# total head count through _SeqAllToAll and its backward pass, because the gather direction
-# cannot recover it from the tensor shape alone.
-_ulysses_num_kv_heads = None
 
+def _resolve_kv_heads(input, scatter_idx, seq_world_size, num_kv_heads):
+    """KV head count this all-to-all partitions against, or None when the caller has none.
 
-def set_ulysses_num_kv_heads(num):
-    global _ulysses_num_kv_heads
-    _ulysses_num_kv_heads = num
-
-
-def get_ulysses_num_kv_heads():
-    return _ulysses_num_kv_heads
+    Only the scatter direction can read a count off the tensor; the gather direction and
+    ``_SeqAllToAll.backward`` see an already-sharded head dim. The value resolved here is
+    therefore carried on the autograd context and replayed, rather than recomputed or kept
+    process-wide where a second model would inherit it.
+    """
+    if num_kv_heads is None and not scatter_idx < 2 and input.shape[2] % seq_world_size != 0:
+        # Scatter direction, so dim 2 is still the full head dim and an uneven count is visible
+        # here. Taking it as the partition basis assumes one KV group per head: under GQA a
+        # query tensor carries a multiple of the KV group count, and splitting on that number
+        # would tear a group across ranks. GQA callers pass the count in instead.
+        num_kv_heads = input.shape[2]
 
+    if num_kv_heads is not None:
+        assert num_kv_heads >= seq_world_size, (
+            f"Number of key-value heads ({num_kv_heads}) must be at least the sequence parallel "
+            f"size ({seq_world_size}); a smaller count leaves a rank with no head to attend over.")
 
-def _ulysses_meta():
-    return AutoTPMeta(num_kv_heads=_ulysses_num_kv_heads)
+    return num_kv_heads
 
 
 try:
@@ -153,13 +154,16 @@ def apply_rotary_pos_emb(t, freqs_cos, freqs_sin):
     return res
 
 
-def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
+def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group, num_kv_heads):
     seq_world_size = dist.get_world_size(group)
+    # AutoTPMeta is built here rather than threaded: the partition follows this all-to-all's own
+    # KV count, and the shard helper is the only thing that needs it in that shape.
+    meta = AutoTPMeta(num_kv_heads=num_kv_heads)
     inp_shape = list(input.shape)
     assert batch_dim_idx in [0, 1], "batch_dim_idx must be either 0 or 1"
 
     if not (scatter_idx < 2):
-        input_splits = get_shard_size_list(inp_shape[scatter_idx], seq_world_size, _ulysses_meta())
+        input_splits = get_shard_size_list(inp_shape[scatter_idx], seq_world_size, meta)
         input = input.transpose(0, scatter_idx).contiguous()
         local_heads = input_splits[groups._get_sequence_parallel_rank()]
         output_splits = [local_heads] * seq_world_size
@@ -193,7 +197,7 @@ def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
         elif batch_dim_idx == 1:  #s,b,h
             input = input.transpose(1, 2).contiguous()  #s,h,b
         seq_len, h, batch_size = input.shape
-        num_local_heads_list = get_shard_size_list(get_ulysses_num_kv_heads(), seq_world_size, _ulysses_meta())
+        num_local_heads_list = get_shard_size_list(num_kv_heads, seq_world_size, meta)
         local_heads = num_local_heads_list[groups._get_sequence_parallel_rank()]
         h_dim = h // local_heads
         local_seq_len = seq_len // seq_world_size
@@ -204,7 +208,7 @@ def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
         coeff = local_seq_len_with_heads // local_heads  #per head: dim size of local_seq_len*hdim
 
         #uneven seq_world_size coeff,  total_heads/local_heads.
-        heads_scale_coeff = get_ulysses_num_kv_heads() / local_heads
+        heads_scale_coeff = num_kv_heads / local_heads
 
         output_splits = [num_local_heads * coeff for num_local_heads in num_local_heads_list]
         output_buff_d1_size = int(heads_scale_coeff * local_seq_len_with_heads)
@@ -221,9 +225,9 @@ def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
         #total_num_large_heads=sum([2,2,2])=7
         #total_num_small_heads=sum([1])=1
 
-        chunk_num_heads_small = get_ulysses_num_kv_heads() // seq_world_size  # even heads compatible
+        chunk_num_heads_small = num_kv_heads // seq_world_size  # even heads compatible
         chunk_num_heads_large = chunk_num_heads_small + 1
-        num_chunk_heads_large = get_ulysses_num_kv_heads() % seq_world_size
+        num_chunk_heads_large = num_kv_heads % seq_world_size
         num_chunk_heads_small = seq_world_size - num_chunk_heads_large
         total_num_large_heads = num_chunk_heads_large * chunk_num_heads_large
         total_num_small_heads = num_chunk_heads_small * chunk_num_heads_small
@@ -263,21 +267,25 @@ def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
     return output
 
 
-def single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, async_op=False, handle=None, type=None):
+def single_all_to_all(input,
+                      scatter_idx,
+                      gather_idx,
+                      batch_dim_idx,
+                      group,
+                      async_op=False,
+                      handle=None,
+                      type=None,
+                      num_kv_heads=None):
     seq_world_size = dist.get_world_size(group)
-    # we only need num_heads once
-    num_heads = input.shape[2]
-
-    if get_ulysses_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and not scatter_idx < 2):
-        # Assuming here that the number of heads for q is consistent with kv
-        # If not, additional logic is required for cases like GQA
-        if get_ulysses_num_kv_heads() is None:
-            assert num_heads > seq_world_size, f"Number of heads ({num_heads}) must be larger than sequence parallel size ({seq_world_size})"
-            # set heads at first call by num_total_heads.
-            # then use ``get_ulysses_num_kv_heads() is not None`` to re-entry uneven path.
-            set_ulysses_num_kv_heads(num_heads)
+
+    num_kv_heads = _resolve_kv_heads(input, scatter_idx, seq_world_size, num_kv_heads)
+
+    # Only an indivisible count needs the uneven implementation. Taking it for every explicit
+    # count would push evenly-split models onto it too, and it rejects async_op, which is how
+    # the overlapped q/k path runs.
+    if num_kv_heads is not None and num_kv_heads % seq_world_size != 0:
         assert async_op == False, "uneven head sp does not support async op"
-        return uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group)
+        return uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group, num_kv_heads)
 
     pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape = _generate_layout_params(
         scatter_idx, batch_dim_idx, seq_world_size, input)
@@ -331,7 +339,8 @@ def forward(ctx: Any,
                 stream=None,
                 handle=None,
                 type=None,
-                is_fwd=True) -> Tensor:
+                is_fwd=True,
+                num_kv_heads=None) -> Tensor:
         ctx.group = group
         ctx.scatter_idx = scatter_idx
         ctx.gather_idx = gather_idx
@@ -339,29 +348,67 @@ def forward(ctx: Any,
         ctx.handle = handle
         ctx.type = type
         ctx.batch_dim_idx = batch_dim_idx
+        # Resolve once, here: backward runs with scatter and gather swapped and cannot recover
+        # the count from its own already-sharded tensor, so it replays this value.
+        num_kv_heads = _resolve_kv_heads(input, scatter_idx, dist.get_world_size(group), num_kv_heads)
+        ctx.num_kv_heads = num_kv_heads
         if ctx.handle is None:
-            res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False)
+            res = single_all_to_all(input,
+                                    scatter_idx,
+                                    gather_idx,
+                                    batch_dim_idx,
+                                    group,
+                                    False,
+                                    num_kv_heads=num_kv_heads)
 
         else:
             # overlap communication path
             if not is_fwd and type == 'o':
                 assert ctx.stream != None
-                res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False)
+                res = single_all_to_all(input,
+                                        scatter_idx,
+                                        gather_idx,
+                                        batch_dim_idx,
+                                        group,
+                                        False,
+                                        num_kv_heads=num_kv_heads)
                 get_accelerator().current_stream().wait_stream(ctx.stream)
                 # The computation of d o_weight can overlap with the communication of d o_input
 
             elif not is_fwd and type in ('q', 'k'):
                 # Achieve communication overlap by pipelining the matrix computation and communication of dq, dk, and dv
                 type = 'd' + type
-                res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, True, handle, type)
+                res = single_all_to_all(input,
+                                        scatter_idx,
+                                        gather_idx,
+                                        batch_dim_idx,
+                                        group,
+                                        True,
+                                        handle,
+                                        type,
+                                        num_kv_heads=num_kv_heads)
 
             elif is_fwd and type in ('q', 'k'):
                 # Achieve communication overlap by pipelining the matrix computation and communication of q, k, and v
                 type = 'fwd_' + type
-                res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False, handle, type)
+                res = single_all_to_all(input,
+                                        scatter_idx,
+                                        gather_idx,
+                                        batch_dim_idx,
+                                        group,
+                                        False,
+                                        handle,
+                                        type,
+                                        num_kv_heads=num_kv_heads)
 
             else:
-                res = single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, False)
+                res = single_all_to_all(input,
+                                        scatter_idx,
+                                        gather_idx,
+                                        batch_dim_idx,
+                                        group,
+                                        False,
+                                        num_kv_heads=num_kv_heads)
 
         return res
 
@@ -370,7 +417,8 @@ def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]:
 
         return (None,
                 _SeqAllToAll.apply(ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx, ctx.batch_dim_idx,
-                                   ctx.stream, ctx.handle, ctx.type, False), None, None, None, None, None, None, None)
+                                   ctx.stream, ctx.handle, ctx.type, False,
+                                   ctx.num_kv_heads), None, None, None, None, None, None, None, None)
 
 
 class DistributedAttention(torch.nn.Module):
@@ -381,6 +429,9 @@ class DistributedAttention(torch.nn.Module):
         sequence_process_group (ProcessGroup): sequence parallel process group
         scatter_idx (int): scatter_idx for all2all comm
         gather_idx (int): gather_idx for all2all comm
+        num_kv_heads (int): number of key-value heads before sharding, which is what the uneven
+            partition follows. Only needed when it does not divide evenly by the sequence
+            parallel world size; when omitted it is read from the key tensor of each forward.
     """
 
     def __init__(
@@ -390,11 +441,15 @@ def __init__(
         scatter_idx: int = 2,
         gather_idx: int = 0,
         sp_stream=None,
+        num_kv_heads: int = None,
     ) -> None:
 
         super(DistributedAttention, self).__init__()
         self.local_attn = local_attention
         self.spg = sequence_process_group
+        # Kept per instance: a second model with a different head count must not decide how this
+        # one is sharded.
+        self.num_kv_heads = num_kv_heads
         self.scatter_idx = scatter_idx
         self.gather_idx = gather_idx
         self.sp_overlap_comm = False
@@ -447,17 +502,22 @@ def pre_hook_fun(grad):
 
             return pre_hook_fun
 
+        # The partition follows KV groups, so the count comes off the key tensor: under GQA a
+        # query head has to land on the rank holding its KV head, and splitting on the query
+        # count would tear the group apart.
+        num_kv_heads = self.num_kv_heads if self.num_kv_heads is not None else key.shape[2]
+
         self.layer_sync(query)
         query_layer = _SeqAllToAll.apply(self.spg, query, self.scatter_idx, self.gather_idx, batch_dim_idx, None,
-                                         self.overlap_handles, 'q')
+                                         self.overlap_handles, 'q', True, num_kv_heads)
         self.layer_sync(key)
         key_layer = _SeqAllToAll.apply(self.spg, key, self.scatter_idx, self.gather_idx, batch_dim_idx, None,
-                                       self.overlap_handles, 'k')
+                                       self.overlap_handles, 'k', True, num_kv_heads)
         if self.sp_overlap_comm:
             self.default_stream.wait_stream(self.sp_stream)
 
         value_layer = _SeqAllToAll.apply(self.spg, value, self.scatter_idx, self.gather_idx, batch_dim_idx, None,
-                                         self.overlap_handles, 'v')
+                                         self.overlap_handles, 'v', True, num_kv_heads)
 
         if self.sp_overlap_comm:
             # Register a hook to synchronize dq and dk after the all-to-all
@@ -479,7 +539,7 @@ def pre_hook_fun(grad):
         context_layer = self.local_attn(query_layer, key_layer, value_layer, *args, **kwargs)
 
         output = _SeqAllToAll.apply(self.spg, context_layer, self.gather_idx, self.scatter_idx, batch_dim_idx,
-                                    self.sp_stream, self.overlap_handles, 'o')
+                                    self.sp_stream, self.overlap_handles, 'o', True, num_kv_heads)
 
         #out e.g., [s/p::h]
         return output
diff --git a/docs/_tutorials/ulysses-alst-sequence-parallelism.md b/docs/_tutorials/ulysses-alst-sequence-parallelism.md
index edce226749c6..5674c0ff047a 100644
--- a/docs/_tutorials/ulysses-alst-sequence-parallelism.md
+++ b/docs/_tutorials/ulysses-alst-sequence-parallelism.md
@@ -226,7 +226,7 @@ In theory you could just average `losses_per_rank`, but the system supports vari
 
 If you are using Sequence Parallelism with **PyTorch version < 2.3**, you may encounter an `IndexError: tuple index out of range` during the backward pass when `sequence_parallel_size < world_size`. This is due to a known issue in the `torch.distributed.all_gather` backward implementation in older versions.
 
-**Workaround:** We recommend using a **weighted `all_reduce` pattern** instead of `all_gather` for loss averaging. You can refer to our [regression test case](https://github.com/deepspeedai/DeepSpeed/blob/master/tests/unit/sequence_parallelism/test_ulysses.py) for a code example of this workaround.
+**Workaround:** We recommend using a **weighted `all_reduce` pattern** instead of `all_gather` for loss averaging. You can refer to our [regression test case](https://github.com/deepspeedai/DeepSpeed/blob/master/tests/unit/v1/sequence_parallelism/test_ulysses.py) for a code example of this workaround.
 
 ### Why do labels need to be pre-shifted?
 
diff --git a/tests/unit/sequence_parallelism/test_autosp_equivalence.py b/tests/unit/v1/sequence_parallelism/test_autosp_equivalence.py
similarity index 99%
rename from tests/unit/sequence_parallelism/test_autosp_equivalence.py
rename to tests/unit/v1/sequence_parallelism/test_autosp_equivalence.py
index 0c289353a29d..fe2c65625579 100644
--- a/tests/unit/sequence_parallelism/test_autosp_equivalence.py
+++ b/tests/unit/v1/sequence_parallelism/test_autosp_equivalence.py
@@ -9,7 +9,7 @@
 These tests require 2 GPUs.
 Run with:
 
-    NCCL_P2P_DISABLE=1 python -m pytest tests/unit/sequence_parallelism/test_autosp_equivalence.py -v
+    NCCL_P2P_DISABLE=1 python -m pytest tests/unit/v1/sequence_parallelism/test_autosp_equivalence.py -v
 """
 
 import torch
diff --git a/tests/unit/sequence_parallelism/test_autosp_integration.py b/tests/unit/v1/sequence_parallelism/test_autosp_integration.py
similarity index 99%
rename from tests/unit/sequence_parallelism/test_autosp_integration.py
rename to tests/unit/v1/sequence_parallelism/test_autosp_integration.py
index 4efcdb07c302..c374efddae75 100644
--- a/tests/unit/sequence_parallelism/test_autosp_integration.py
+++ b/tests/unit/v1/sequence_parallelism/test_autosp_integration.py
@@ -15,7 +15,7 @@
 These tests require 2 GPUs.
 Run with:
 
-    NCCL_P2P_DISABLE=1 python -m pytest tests/unit/sequence_parallelism/test_autosp_integration.py -v
+    NCCL_P2P_DISABLE=1 python -m pytest tests/unit/v1/sequence_parallelism/test_autosp_integration.py -v
 """
 
 import torch
diff --git a/tests/unit/sequence_parallelism/test_ulysses.py b/tests/unit/v1/sequence_parallelism/test_ulysses.py
similarity index 77%
rename from tests/unit/sequence_parallelism/test_ulysses.py
rename to tests/unit/v1/sequence_parallelism/test_ulysses.py
index 14a54fe88a20..4a77cc35e1e8 100644
--- a/tests/unit/sequence_parallelism/test_ulysses.py
+++ b/tests/unit/v1/sequence_parallelism/test_ulysses.py
@@ -7,11 +7,13 @@
 import torch
 import torch.nn.functional as F
 import deepspeed.comm as dist
+from deepspeed.accelerator import get_accelerator
 from deepspeed import initialize
 import deepspeed.runtime.sequence_parallel.parallel_state_sp as sp_mpu
 from transformers import AutoModel
 from unit.common import DistributedTest
-from deepspeed.sequence.layer import _SeqAllToAll, _generate_layout_params, post_all2all, pre_all2all_fun
+from deepspeed.sequence.layer import (DistributedAttention, _SeqAllToAll, _generate_layout_params, post_all2all,
+                                      pre_all2all_fun, single_all_to_all)
 from deepspeed.sequence.fpdt_layer import _FPDTGPUOffloadingAttentionImpl_, FPDT_InputConstruct
 from unit.util import skip_on_arch
 from unit.simple_model import *
@@ -256,7 +258,7 @@ def seq_batch_heads_hash(d0, d1, h, offset_d0=0, offset_d1=0, offset_h=0):
 
             ### first all2all: sequence parallel to head parallel
             s2h_tensor = _SeqAllToAll.apply(ds_engine.seq_parallel_group, input_tensor, scatter_idx, gather_idx,
-                                            batch_dim_idx)
+                                            batch_dim_idx, None, None, None, True, num_heads)
 
             # s2h_tensor check for the first all2all: compare with the expected ground truth
             d0_indices = torch.arange(s2h_tensor.shape[0]).reshape(-1, 1, 1, 1)
@@ -271,8 +273,10 @@ def seq_batch_heads_hash(d0, d1, h, offset_d0=0, offset_d1=0, offset_h=0):
                                   s2h_tensor), f"s2h_tensor differs from the expected for sequence dim: {seq_dim}"
             #No op
             ### second all2all: head parallel to sequence parallel
+            # The gather direction cannot read an uneven count off an already-sharded tensor,
+            # so this call site supplies it too.
             h2s_tensor = _SeqAllToAll.apply(ds_engine.seq_parallel_group, s2h_tensor, gather_idx, scatter_idx,
-                                            batch_dim_idx)
+                                            batch_dim_idx, None, None, None, True, num_heads)
             print(
                 f'[{dist.get_rank()}] s={seq_dim} input: {input_tensor.shape} s2h: {s2h_tensor.shape} h2s_tensor: {h2s_tensor.shape}'
             )
@@ -433,3 +437,115 @@ def test_sp_loss_backward_stability(self, sp_size: int) -> None:
         expected_grad = 0.5
         assert torch.allclose(local_loss.grad, torch.tensor(expected_grad, device=ds_engine.device)), \
             f"Gradient mismatch! Expected {expected_grad}, got {local_loss.grad}"
+
+
+class _RecordingAttention(torch.nn.Module):
+    """Stands in for the local attention and records the head shard it was handed."""
+
+    def __init__(self):
+        super().__init__()
+        self.head_ids = None
+
+    def forward(self, query, key, value, *args, **kwargs):
+        # Each head carries its own global index, so the shard is readable off the tensor.
+        self.head_ids = sorted({int(v) for v in query[0, 0, :, 0].tolist()})
+        # Fold the query heads onto the kv width the way GQA groups them, so the output is the
+        # shape the reverse all-to-all expects and all three inputs carry a gradient.
+        group_size = query.shape[2] // value.shape[2]
+        grouped_query = query.view(*query.shape[:2], value.shape[2], group_size, query.shape[3]).mean(dim=3)
+        return grouped_query + key + value
+
+
+class TestUlyssesKVHeadCount(DistributedTest):
+    """The count the uneven all-to-all splits against is threaded per call (#8291).
+
+    It used to sit in a process-wide slot memoized on the first uneven all-to-all, so the first
+    model to take that path decided how every later one was sharded.
+    """
+    world_size = 2
+    LOCAL_SEQ = 4
+    HEAD_DIM = 8
+
+    def _sequence_parallel_group(self):
+        groups.mesh_device = dist.initialize_mesh_device((1, self.world_size), ("data_parallel", "sequence_parallel"))
+        return groups.mesh_device.get_group(mesh_dim="sequence_parallel")
+
+    def _tagged(self, num_heads, batch_dim_idx=0):
+        device = get_accelerator().current_device_name()
+        shape = (1, self.LOCAL_SEQ) if batch_dim_idx == 0 else (self.LOCAL_SEQ, 1)
+        tensor = torch.zeros(*shape, num_heads, self.HEAD_DIM, device=device, requires_grad=True)
+        with torch.no_grad():
+            tensor[:] = torch.arange(num_heads, device=device).view(1, 1, -1, 1).float()
+        return tensor
+
+    def _run(self, attn, num_heads, num_kv_heads=None):
+        """One forward + backward; returns the query head ids this rank received."""
+        query = self._tagged(num_heads)
+        key = self._tagged(num_kv_heads if num_kv_heads is not None else num_heads)
+        output = attn(query, key, key.clone(), 0)
+        output.sum().backward()
+        return attn.local_attn.head_ids
+
+    def test_second_model_does_not_reshard_the_first(self):
+        # 3 heads over 2 ranks is uneven ([2, 1]) and so is 5 ([3, 2]). The two splits are
+        # distinguishable, so a shared count shows up as one model taking the other's.
+        sp_group = self._sequence_parallel_group()
+        rank = dist.get_rank(group=sp_group)
+
+        teacher = DistributedAttention(_RecordingAttention(), sp_group, scatter_idx=2, gather_idx=1)
+        student = DistributedAttention(_RecordingAttention(), sp_group, scatter_idx=2, gather_idx=1)
+        expected_teacher = [[0, 1], [2]][rank]
+
+        assert self._run(teacher, 3) == expected_teacher
+        assert self._run(student, 5) == [[0, 1, 2], [3, 4]][rank]
+        assert self._run(teacher, 3) == expected_teacher
+
+    @pytest.mark.parametrize("batch_dim_idx", [0, 1])
+    def test_gqa_partitions_by_kv_groups(self, batch_dim_idx):
+        # Q=6 / KV=3 over 2 ranks. 6 divides evenly, but a query head has to stay on the rank
+        # holding its KV head, so the split follows the KV groups [2, 1] and Q becomes [4, 2].
+        sp_group = self._sequence_parallel_group()
+        rank = dist.get_rank(group=sp_group)
+
+        gather_idx = 1 if batch_dim_idx == 0 else 0
+        attn = DistributedAttention(_RecordingAttention(), sp_group, scatter_idx=2, gather_idx=gather_idx)
+        query = self._tagged(6, batch_dim_idx)
+        key = self._tagged(3, batch_dim_idx)
+        value = self._tagged(3, batch_dim_idx)
+
+        output = attn(query, key, value, batch_dim_idx)
+        output.sum().backward()
+
+        assert attn.local_attn.head_ids == [[0, 1, 2, 3], [4, 5]][rank]
+        for tensor in (query, key, value):
+            assert tensor.grad is not None and torch.isfinite(tensor.grad).all()
+
+    def test_fewer_kv_heads_than_ranks_is_rejected_before_the_collective(self):
+        # 1 KV head over 2 ranks leaves rank 1 with nothing to attend over. Both ranks have to
+        # reject it together: one of them raising inside the all-to-all hangs the other.
+        sp_group = self._sequence_parallel_group()
+
+        attn = DistributedAttention(_RecordingAttention(), sp_group, scatter_idx=2, gather_idx=1)
+        query = self._tagged(2)
+        key = self._tagged(1)
+
+        with pytest.raises(AssertionError, match="at least the sequence parallel size"):
+            attn(query, key, key.clone(), 0)
+
+    def test_even_count_keeps_the_fast_path_and_its_async_op(self):
+        # An explicit count that divides the world size must not be routed to the uneven
+        # implementation, which rejects async_op and so would disable the overlapped q/k path.
+        sp_group = self._sequence_parallel_group()
+        handle = {}
+
+        output = single_all_to_all(self._tagged(4), 2, 1, 0, sp_group, True, handle, 'dq', num_kv_heads=4)
+
+        assert output.shape[2] == 4 // self.world_size
+
+    def test_uneven_count_still_takes_the_uneven_path(self):
+        sp_group = self._sequence_parallel_group()
+        rank = dist.get_rank(group=sp_group)
+
+        output = single_all_to_all(self._tagged(3), 2, 1, 0, sp_group, num_kv_heads=3)
+
+        assert output.shape[2] == [2, 1][rank]