[megatron] KDA context parallelism: opt-in all-gather exchange (SKYRL_KDA_CP_EXCHANGE=allgather)

[avigyabb]

Stacked on #2418 (KDA context parallelism, all-to-all exchange). Review only the top commit.

Summary

Under context parallelism, KDA moves from sequence to head parallelism. #2418 does that with megatron-core's GatedDeltaNet all-to-all of the projected q/k/v/f/g/beta, which are ~1.25x wider than the hidden state and run on TP-gathered rows.

SKYRL_KDA_CP_EXCHANGE=allgather (opt-in) instead all-gathers the sequence-parallel hidden-state shards over CP (the only cross-node hop when CP spans nodes), then over TP, and projects only this rank's 1/cp slice of the local heads. One permutation of the narrow projected tensors puts tokens into natural order; backward reduce-scatters in reverse. About 5x fewer bytes cross the CP group.

Performance

1M tokens/seq, TP8 x CP2, EP32 x ETP2, 64 B200	all-to-all (#2418)	all-gather (this PR)
NCCL over TCP	fwd+bwd 1482 s	1251 s
NCCL over EFA	335 s/step, 12.5k tok/s	307 s/step, 13.7k tok/s (−8%)

Over EFA the step is 8% faster at the same peak memory (144.6 vs 144.5 GiB); EFA rows are the mean warm step of the full step (fwd+bwd+optimizer) with the rest of the long-context recipe on. On a 4-layer GLM-5.3-Flash slice at 1M tokens, CP2 across 2 nodes over EFA: 34.9 -> 32.8 s/step (−6%), same memory.

Testing

• gpu_ci/patches/megatron/mcore_ext/test_kda_context_parallel.py, parametrized over both exchanges: CP=2 outputs and input/param grads match CP=1. Passed on B200 on this exact code before the split (as part of [megatron] Context parallelism for GLM-5.3-Flash (head-wise KDA, all-gather DSA) #2418).
• Over-EFA 1M-token comparison against [megatron] Context parallelism for GLM-5.3-Flash (head-wise KDA, all-gather DSA) #2418's all-to-all (full model and 4-layer slice, numbers above).

Split out of #2418.

🤖 Generated with Claude Code

[avigyabb]

all gather was better in all tests - maybe we should fold into one path? maybe just wait until megatron PR gets pushed?