Fix the seq-first Ulysses all2all output layout

[vineethsaivs]

Symptom

_generate_layout_params (deepspeed/sequence/layer.py) builds the reshape target for every all2all in DistributedAttention. For the seq-first layout (batch_dim_idx=1, (s, b, n, h)) with scatter_idx < 2 it returns

post_all2all_res_shape = [bs, seq_world_size * global_seq_len, num_local_head // seq_world_size, head_dim]

That is the batch_dim_idx=0 / scatter_idx >= 2 shape: batch first, sequence multiplied, heads divided. This direction does the opposite, it scatters the sequence and gathers the heads, so the result should be [global_seq_len // seq_world_size, bs, seq_world_size * num_local_head, head_dim].

The element count still matches whenever num_local_head is divisible by seq_world_size, so the reshape succeeds and silently returns a transposed, mis-strided tensor. When it is not divisible the floor division makes a dimension 0 and the reshape raises shape '[...]' is invalid for input of size ....

Root cause

A copy of the wrong sibling branch during a refactor. Before #6750 extracted this function, post_all2all computed the shape inline and this case read:

        else:
            # s, b, n, h
            if scatter_idx < 2:
                output = input.permute(1, 2, 0, 3, 4).contiguous()
                output = output.reshape(seq_len // seq_world_size, bs, seq_world_size * num_head,
                                        head_dim).contiguous()

The permute survived the refactor unchanged; only the reshape target was replaced, with the shape from the batch_dim_idx=0/scatter_idx >= 2 branch.

Reachability

DistributedAttention.__init__ defaults to gather_idx=0, and the output-projection all2all swaps the two indices:

output = _SeqAllToAll.apply(self.spg, context_layer, self.gather_idx, self.scatter_idx, batch_dim_idx, ...)

so it runs with scatter_idx=0. _SeqAllToAll.backward swaps them too, so the backward of the q/k/v all2alls takes the same path. Both are the seq-first (s, b, n, h) layout used by Megatron-DeepSpeed.

Why the existing tests miss it

• TestUlyssesAll2All only runs batch_dim_idx = 0.
• TestUlyssesAll2All_odd does cover batch_dim_idx = 1, but its first call has num_heads % seq_world_size != 0, which calls set_num_kv_heads(...). From then on get_num_kv_heads() is not None routes every later call to uneven_heads_all2all, so _generate_layout_params is never reached.
• Both are DistributedTest classes behind skip_on_arch(min_arch=8), so neither runs in the CPU CI.

Fix

One line, restoring the pre-#6750 shape, plus a two-line comment naming the direction.

Test

_generate_layout_params, pre_all2all_fun and post_all2all are pure, so TestUlyssesAll2AllLayout drives them with an emulated all_to_all_single (rank i sends chunk j of dim 0 to rank j) and checks that both directions land the right (sequence, head) shard of a known tensor. No process group, no accelerator, so it runs in the CPU CI, where this branch currently has no coverage at all.

# before, tests/
$ pytest unit/sequence_parallelism/test_ulysses.py -k Layout
FAILED TestUlyssesAll2AllLayout::test_head_to_seq_parallel[2-1] - AssertionError: rank 0 got torch.Size([2, 12, 1, 4]), expected torch.Size([3, 2, 4, 4])
FAILED TestUlyssesAll2AllLayout::test_head_to_seq_parallel[4-1] - RuntimeError: shape '[2, 48, 0, 4]' is invalid for input of size 192
2 failed, 6 passed

# after
8 passed

# whole file, after
8 passed, 43 skipped        (the 43 are the DistributedTest classes, no GPU here)

yapf --style .style.yapf and flake8 are clean on both files.

[vineethsaivs]

Flagging some prior art I only found today: #7868 proposes the same one-line change to post_all2all_res_shape, as part of a larger draft that also pads uneven sequence lengths and touches fpdt_layer.py, sequence/__init__.py and utils/groups.py. That draft has been conflicting and unchanged since 2026-02-22, and the issue it was filed against (#7842) was closed out by its reporter after the mismatch turned out to be their baseline rather than DeepSpeed.

I did not know about it when I opened this. Since the shape bug is real independently of the padding work, this PR keeps to just that one line plus a CPU regression test that fails on master (2 failed, 6 passed) and passes with the fix (8 passed). Happy to close this in favour of #7868 if a maintainer would rather revive the larger change.