[Ulysses] Carry the KV head count per DistributedAttention

[alanhuangyoo]

Fixes #8291.

What breaks

single_all_to_all decides whether a model does uneven-head sequence parallelism by asking whether the process-wide tp_shard head count is set:

if get_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and not scatter_idx < 2):

AutoTP writes that same slot on every injection (replace_module.py:306, engine.py:770), so it is not only a second Ulysses model that clobbers the first — a single AutoTP model anywhere in the process is enough. uneven_heads_all2all then splits against that value, and get_shard_size_list falls back to the same global internally, so threading a value without also passing num_kv_heads= explicitly would not have been enough.

A 12-head model on sp=4 should hand every rank 3 heads. With a leftover num_kv_heads=3:

                              12 heads, sp=4
  clean process            -> [3, 3, 3, 3]
  after AutoTP set kv=3    -> [4, 4, 4, 0]

Driving it through DistributedAttention on 4 gloo ranks, upstream 6e3bd087f:

clean process
  rank0: local_attn got 3 heads  ids=[0, 1, 2]     backward OK
  rank1: local_attn got 3 heads  ids=[3, 4, 5]     backward OK
  rank2: local_attn got 3 heads  ids=[6, 7, 8]     backward OK
  rank3: local_attn got 3 heads  ids=[9, 10, 11]   backward OK

after an AutoTP model ran first
  (no output — rank 3 gets an empty tensor and the forward hangs in the output all-to-all)

Same script with this branch applied, both cases:

  rank0: local_attn got 3 heads  ids=[0, 1, 2]     backward OK
  rank1: local_attn got 3 heads  ids=[3, 4, 5]     backward OK
  rank2: local_attn got 3 heads  ids=[6, 7, 8]     backward OK
  rank3: local_attn got 3 heads  ids=[9, 10, 11]   backward OK

The change

DistributedAttention holds the count, either from a new num_kv_heads argument or read off its own key tensor, and only when it does not divide evenly — claiming one otherwise would push every model onto the uneven kernels. It threads the value through _SeqAllToAll into single_all_to_all and uneven_heads_all2all, and backward restores it from ctx, which is the piece the gather direction cannot recover from tensor shapes.

The value is the KV head count, taken from the key tensor rather than the query tensor. That is what the tp_shard slot held, and it is what the partition has to follow: under GQA a query head must land on the rank holding its KV head. Q=6 / KV=3 over sp=2 divides evenly on the query side, so reading the query would give [3, 3] and split group 1 across two ranks; the KV side gives [2, 1] and therefore [4, 2]. Thanks to @FU-max-boop for catching that on the issue - my first revision had exactly this bug.

get_shard_size_list is called with num_kv_heads set explicitly so its module-level default cannot reintroduce the leak.

Callers that reach single_all_to_all directly, Megatron-DeepSpeed among them, pass nothing and keep the old behaviour including the set_num_kv_heads publish on first use. A sentinel separates "not supplied" from a deliberate None, which now means "splits evenly, take the fast path".

This is independent of #8241 — it touches only deepspeed/sequence/layer.py and does not depend on AutoTPMeta — but the two overlap in intent, and I am happy to rebase on it if that lands first.

Tests

New tests/unit/sequence_parallelism/test_ulysses_multiple_models.py, shaped after the multi-model AutoTP test in #8241:

• two DistributedAttention instances at 3 and 5 heads over sp=2 — the first still splits [2, 1] after the second has run forward and backward
• a 6-head model over sp=2 with a leftover num_kv_heads=3 — this is the case above, and it hangs on master
• GQA, Q=6 / KV=3 over sp=2 - splits [4, 2], matching master
• the explicit num_kv_heads argument

$ LOCAL_SIZE=2 pytest unit/sequence_parallelism/test_ulysses_multiple_models.py
4 passed in 21.65s

$ LOCAL_SIZE=2 pytest unit/sequence_parallelism/ unit/module_inject/test_tp_shard.py
31 passed in 97.39s

Run on CPU, so cpu-torch-latest picks them up rather than only the GPU workflows.

[FU-max-boop]

Thanks for taking the dependency-free route, and for incorporating the GQA partitioning point. I ran exact head b4f148056b61586fcdc8970f20504ba93f5c075d against base 6e3bd087f9bee031a6853891e35ead44a80122ad with a deterministic 2-rank CPU/Gloo harness. I found two regressions that look blocking before this version lands.

1. Legacy direct _SeqAllToAll.apply loses the inferred odd head count before backward

When the caller omits num_kv_heads, _SeqAllToAll.forward resolves _KV_HEADS_FROM_GLOBAL to the current global value (None) before calling single_all_to_all. The scatter-side helper can then infer num_heads == 3, but it no longer knows the sentinel was used, so the resolved count is stored in neither ctx.num_kv_heads nor the compatibility global. Backward receives explicit None; with scatter/gather swapped, it selects the equal-split path and the ranks disagree on collective sizes.

With a direct odd-head forward/backward (heads=3, SP=2, requires_grad=True):

base:
(1, 'ok', (1, 8, 1, 2), True)
(0, 'ok', (1, 8, 2, 2), True)
('exitcodes', [0, 0])

#8316:
gloo::EnforceNotMet: op.nread == op.preamble.nbytes
(0, 'hung', ...)
('exitcodes', [-15, -6])

A minimal direction would be to preserve whether the argument was unspecified, resolve/infer the effective count once on the scatter side, store that effective value in the autograd context, and replay the same value in backward. A direct _SeqAllToAll.apply odd-head autograd regression would cover the compatibility path.

2. num_kv_heads < world_size now enters a collective with a zero-head rank

For Q=2 / KV=1 / SP=2, an explicit or key-inferred num_kv_heads=1 makes the new num_kv_heads is not None branch bypass the old pre-collective head-count guard. Rank 1 receives local_heads == 0, reaches h_dim = h // local_heads, and raises while rank 0 is already waiting in a collective.

base:
rank 0: AssertionError: Number of heads (1) must be larger than sequence parallel size (2)
rank 1: AssertionError: Number of heads (1) must be larger than sequence parallel size (2)
('exitcodes', [0, 0])  # worker reports the caught error, then exits normally

#8316:
(0, 'hung', ...)
(1, 'err', 'ZeroDivisionError', 'integer division or modulo by zero', ...)
deepspeed/sequence/layer.py:183: h_dim = h // local_heads
('exitcodes', [-15, 0])

This needs a rank-consistent validation before the first collective: reject num_kv_heads < world_size (or define a supported zero-head protocol) on every rank. A Q=2 / KV=1 / SP=2 regression should assert synchronized fail-fast behavior.

Reproducer: https://gist.github.com/FU-max-boop/2da2d1c018af3da35e24f04829da1445
SHA-256: 22c4359ce38b92a76737c378c97b5fab73d2dda22881ca4e79b83963ef383120

Run from either checkout root:

PYTHONPATH="$PWD" DS_ACCELERATOR=cpu python /path/to/repro_8316_p1.py direct-odd --timeout 12
PYTHONPATH="$PWD" DS_ACCELERATOR=cpu python /path/to/repro_8316_p1.py mqa-kv-lt-world --timeout 12

The current Modal job is green, but its execution log is 79 passed, 9 skipped from the compile suite; it did not exercise the new Ulysses test file or the existing direct odd-head regression, so it does not cover these cases.

The dependency-free approach is the preferable landing path if these are fixed. I am happy to rerun the harness and review the next head. I will keep my stacked total_heads / partition_heads implementation only as a fallback rather than opening a competing PR while this one is active.

[alanhuangyoo]

Both reproduce here, thanks — the harness made them easy to confirm.

Same two cases on a 2-rank gloo harness, against 6e3bd087f and against b4f148056:

upstream master
  direct apply, heads=3, SP=2   rank0: ok  out=(1, 8, 2, 2)  grad=True
                                rank1: ok  out=(1, 8, 1, 2)  grad=True
  Q=2 / KV=1 / SP=2             both ranks: AssertionError: Number of heads (1) must be
                                larger than sequence parallel size (2)

b4f148056 (the version you reviewed)
  direct apply, heads=3, SP=2   no output, both ranks hung
  Q=2 / KV=1 / SP=2             no output, both ranks hung

Your reading of the first one is right. _SeqAllToAll.forward resolved the sentinel to None before calling single_all_to_all, so reads_global came out false there, and the count single_all_to_all then inferred from the tensor went into neither ctx nor the compatibility global. Backward got an explicit None, and with scatter and gather swapped that is the even path.

The second one is the num_heads > seq_world_size assertion. It only ever ran on the lazy branch, so threading a count walked straight past it and rank 1 reached h // 0 while rank 0 was already inside the collective.

Both come from resolution being split across two places, so it now happens in one, _resolve_kv_heads, called by _SeqAllToAll.forward before it stores ctx.num_kv_heads and by single_all_to_all for callers that arrive there directly. It reads the global, infers from the tensor on the scatter side, publishes back to the global when the caller passed nothing, and asserts once for everyone before any collective. The guard is now >= seq_world_size rather than >: the old form was only reachable when the count did not divide the world size, where the two agree, and > would reject a legitimate explicit num_kv_heads == sp_size.

After that, both cases match master exactly:

  direct apply, heads=3, SP=2   rank0: ok  out=(1, 8, 2, 2)  grad=True
                                rank1: ok  out=(1, 8, 1, 2)  grad=True
  Q=2 / KV=1 / SP=2             both ranks: AssertionError: Number of key-value heads (1)
                                must be at least the sequence parallel size (2)

Added both as regression tests: test_direct_all_to_all_replays_the_inferred_count_in_backward covers the compatibility path including backward and the global publish, and test_fewer_kv_heads_than_ranks_is_rejected_before_the_collective pins the guard. Six tests in the file now, and unit/sequence_parallelism/ plus unit/module_inject/test_tp_shard.py stay green.

Pushed. Worth another look when you have time — and the offer stands, if you would rather land your stacked version on top of #8241 I will close this one.

[delock]

Instead of preserving the global for legacy case, we need to remove the global and fix all call site. Both FPDT and Megatron-DeepSpeed has this information and can call this function with explicit num_kv_heads.

[FU-max-boop]

I audited the current exact head and the relevant in-tree/external callers against this direction.

The smallest global-free route I see is:

• The only in-tree direct callers of single_all_to_all outside layer.py are the 17 FPDT call sites in deepspeed/sequence/fpdt_layer.py. FPDT already carries kv_projection_size and hidden_size_per_attention_head, so its effective KV-head count is kv_projection_size // hidden_size_per_attention_head. That count can be stored on each custom autograd context and threaded explicitly through Q/K/V/output and all reverse-direction calls. The reverse calls cannot reliably infer it from their already-sharded tensor.
• At Megatron-DeepSpeed head aab2f3127c9a5375019221c3a5405ea4cdf98b5e, ParallelAttention already computes self.num_key_value_heads_per_partition in megatron/model/transformer.py:566-567; this local TP-partition count (rather than the model-global config value) is the count available to pass to DistributedAttention(..., num_kv_heads=...) at lines 619-622.
• The Megatron-DeepSpeed FPDT factory also has config.num_key_value_heads when constructing FPDT_Attention (lines 922-971), while the in-tree FPDT implementation can derive the same count from its existing projection metadata.
• The remaining in-tree DistributedAttention constructors are docs/blog examples, the compile test, and this PR's tests. The implicit direct-call regression at test_ulysses_multiple_models.py:111-122 should become an explicit-count backward replay test; the global fixture/imports can then disappear entirely.

If that is the intended contract, I am happy to validate the next DeepSpeed head and/or prepare a small companion Megatron-DeepSpeed caller update after this PR API shape is settled. I will not duplicate the main implementation here.

[FU-max-boop]

Re-reviewed exact head 1d4420b3cc3e360da57c8463590f445ec81ae9bb.

Both previously reported distributed blockers are fixed in the independent 2-rank CPU/Gloo reproductions: the direct odd-head _SeqAllToAll forward/backward now completes with the expected per-rank shapes and gradients, and Q=2 / KV=1 / SP=2 now fails consistently on both ranks before any collective.

I also exercised a Q-width GQA path (Q=6 / KV=3 / SP=2) whose local attention output depends on q, k, and v. For both batch_dim_idx=0 and 1, the full Q/K/V/O forward/backward round trip produced the exact output and expected finite gradients (dq=1, dk=0.5, dv=1) without a collective mismatch or hang.

The repaired delta introduces no remaining P0/P1 finding in my review. The hosted CPU job is green but skips the 2-rank Ulysses cases, so the statements above are based on the exact-head distributed harness rather than inferred from that hosted result. I support landing the dependency-free implementation.

[delock]

here the assumption is non GQA and attention dim is 2. We should put this in comments for better understanding.

[delock]

Hi @alanhuangyoo thanks for your PR. I have left some initial comments. Also this PR needs to be done on top of #8241 ,can you rebase with #8241 ? Thanks!

[FU-max-boop]

I replayed the three commits in 6e3bd087f..1d4420b3c onto exact #8241 head
19e2735099408061716eb7758bc06b83d298f7c8 to check the requested stack. I did not modify the author branch.

The textual rebase is fairly small: the only UU file is deepspeed/sequence/layer.py; the new test is a clean add,
and the next two commits replay cleanly after resolving that first commit. The semantic resolution needs more care,
though:

• [AutoTP] Replace tp_shard process-wide globals with per-model AutoTPMeta #8241 removes get_num_kv_heads / set_num_kv_heads and changes the shard helper to
  get_shard_size_list(total_size, mp_size, meta). Taking the [Ulysses] Carry the KV head count per DistributedAttention #8316 side verbatim therefore leaves an import error and
  num_kv_heads= keyword errors. A narrow fit with [AutoTP] Replace tp_shard process-wide globals with per-model AutoTPMeta #8241 is to keep an integer num_kv_heads in the Ulysses/autograd
  API and construct AutoTPMeta(num_kv_heads=num_kv_heads) only at the shard-helper boundary.
• The fast/uneven choice must be based on num_kv_heads % sequence_parallel_world_size != 0, not merely
  num_kv_heads is not None. Otherwise every explicit even count is forced through the uneven implementation, whose
  async path is intentionally rejected; that would regress ordinary overlap-enabled models.
• Per the review direction, the lower primitive should not read or publish a process global. Resolve the effective
  pre-SP, post-TP count once, save it on _SeqAllToAll's ctx, and replay that integer in backward. Keep the existing
  rank-consistent num_kv_heads >= sp_size validation before the first collective.

The in-tree callers also need to move with the stack. FPDT has 17 direct single_all_to_all calls across its normal
and offload paths; both custom autograd functions already have kv_projection_size and
hidden_size_per_attention_head, so they can derive and save
kv_projection_size // hidden_size_per_attention_head. For Megatron-DeepSpeed, the count to pass at the
DistributedAttention construction is the TP-local self.num_key_value_heads_per_partition, not the model-global
config value.

One CI detail is worth addressing in the same rebase: nv-flash-attn.yml runs test_ulysses.py, not the new
test_ulysses_multiple_models.py; the current hosted green skipped the six new distributed cases. Consolidating the
new regressions into the existing test file (or explicitly updating that workflow) would make the checks meaningful.
On the stacked head I would cover both the even fast/async path and the uneven path, plus Q=6/KV=3/SP=2 Q/K/V/O
forward/backward, KV=1/SP=2 synchronized failure, and the two-model 3->5->3 state-isolation sequence.

[alanhuangyoo]

Rebased onto #8241 at 19e2735099408061716eb7758bc06b83d298f7c8 and reworked against both sets of comments.

@delock on the global: _ulysses_num_kv_heads, its two accessors and _ulysses_meta() are gone, and every in-tree call site now passes the count. FPDT derives kv_projection_size // hidden_size_per_attention_head and saves it on both custom autograd contexts, which covers its 17 direct single_all_to_all calls across the normal and offload paths. AutoTPMeta is constructed only where get_shard_size_list needs it.

One call site that was not on the list: TestUlyssesAll2All_odd does two independent _SeqAllToAll.apply calls rather than a forward/backward pair, and the second runs in the gather direction. It had been reading the count back out of the global, so it now passes it too — two lines.

On the second comment, the shape assumption is written down where the count is inferred: dim 2 being the head dim, and the inference holding only when there is one KV group per head, since a GQA query tensor carries a multiple of the KV count and splitting on that number would tear a group across ranks.

@FU-max-boop the fast/uneven point was a real bug in what I had, not a style note. Threading a count into FPDT's 17 calls plus an explicit constructor argument meant num_kv_heads is not None sent evenly-split models down the uneven path, and that path rejects async_op, so the overlapped q/k calls would have broken. It now keys on num_kv_heads % seq_world_size != 0. Your reading of the #8241 fit was right on the other two as well: integer in the Ulysses/autograd API with AutoTPMeta only at the shard boundary, and the >= sp_size check kept ahead of the first collective.

Tests are consolidated into test_ulysses.py, which nv-flash-attn and hpu-gaudi2-nightly already run, rather than a new file neither picks up. TestUlyssesKVHeadCount covers the even fast path with async_op=True, the uneven path, GQA Q=6/KV=3/SP=2 with gradients on q, k and v, the synchronized KV=1/SP=2 rejection, and the 3 -> 5 -> 3 two-model sequence.

Results:

4x H20, world_size=4
  unit/sequence_parallelism/test_ulysses.py          32 passed, 16 skipped

2 ranks, gloo, CPU
  unit/v1/autotp + unit/module_inject                76 passed, 46 skipped
  TestUlyssesKVHeadCount                             5 passed

yapf / flake8 / check-license                        clean

Megatron-DeepSpeed passing num_key_value_heads_per_partition at the DistributedAttention construction is a separate repo, so it is not in this change. If you want to send that companion, please do.

[alanhuangyoo]

Added batch_dim_idx=1 to the GQA case — the seq-first layout takes a different branch in both _generate_layout_params and uneven_heads_all2all, and the existing odd-head test already covers both, so that one should too. 6 passed on 4x H20.

One thing I could not close, and it is your call @delock. Consolidating into test_ulysses.py puts the regressions in front of nv-flash-attn and hpu-gaudi2-nightly, but neither runs on this PR: nv-flash-attn's last run on this repo was 2026-01-28, and the Required modal-torch-latest selects only under tests/unit/v1. So the checks here stay green without executing any of the six.

Your own #8241 commit says the same thing about tests/unit/model_parallelism — multi-rank DistributedTest that PR CI never ran, fixed by moving it under tests/unit/v1. TestUlyssesKVHeadCount is world_size 2 and has the same problem. Happy to move it to tests/unit/v1/sequence_parallelism/ so the modal job picks it up, or to leave it here if you would rather keep the sequence-parallel tests together. Say which and I will push it.

Everything reported above was run directly rather than inferred from CI: 4x H20 world_size 4 for test_ulysses.py (32 passed, 16 skipped), and 2-rank gloo for unit/v1/autotp plus unit/module_inject (76 passed, 46 skipped).

[delock]

Hi @alanhuangyoo I think you need to move sequence parallel test as well because through CPU we can't get number of device > 1. I'll check this PR after #8241 merged, thanks for your patience. Ping me if it takes too long.

[alanhuangyoo]

Moved, thanks — you are right that CPU cannot give the multi-rank case.

tests/unit/sequence_parallelism/ is now tests/unit/v1/sequence_parallelism/, following what #8241 did for model_parallelism: plain renames plus the one workflow that pointed at the old path (nv-flash-attn.yml, both the paths filter and the pytest target) and the docs link in engine.py. git status shows the three files as renames, so the diff stays readable.

4x H20, world_size=4
  unit/v1/sequence_parallelism/test_ulysses.py    33 passed, 16 skipped

No rush on the review — ping me if anything needs changing once #8241 lands.