Add Packing Support for Context Parallelism (Ring Attention)#2906
Merged
copybara-service[bot] merged 1 commit intoAI-Hypercomputer:mainfrom Apr 20, 2026
Merged
Conversation
kocchop
requested review from
A9isha,
NicoGrande,
NuojCheng,
RissyRan,
SurbhiJainUSC,
aireenmei,
bvandermoon,
gagika,
gobbleturk,
hengtaoguo,
jiangjy1982,
khatwanimohit,
parambole,
richjames0,
shralex,
shuningjin,
suexu1025 and
vipannalla
as code owners
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
richjames0
approved these changes
Dec 31, 2025
Collaborator
richjames0
left a comment
richjames0
left a comment
There was a problem hiding this comment.
a couple of nits but lgtm
gobbleturk
reviewed
Dec 31, 2025
gobbleturk
reviewed
Jan 5, 2026
gobbleturk
approved these changes
Jan 5, 2026
Collaborator
gobbleturk
left a comment
gobbleturk
left a comment
There was a problem hiding this comment.
Thanks for the tests and great comments illustrating the two reorder strategies!
RissyRan
reviewed
Jan 5, 2026
Collaborator
RissyRan
left a comment
RissyRan
left a comment
There was a problem hiding this comment.
LGTM, just minor comments.
RissyRan
approved these changes
Jan 21, 2026
|
This PR has been automatically marked as stale because it has not had recent activity. It will be closed soon if no further activity occurs. Thank you for your contributions. |
|
This PR was closed because it has been inactive for a while. Please reopen it if you are still working on it. |
kocchop
requested review from
dipannita08,
igorts-git and
jesselu-google
as code owners
kocchop
force-pushed
the
faysal/add-thd-2-ring-attn
branch
from
9d04e5a to
e9d76dc
Compare
kocchop
added
enhancement
|
This PR has been automatically marked as stale because it has not had recent activity. It will be closed soon if no further activity occurs. Thank you for your contributions. |
|
This PR was closed because it has been inactive for a while. Please reopen it if you are still working on it. |
kocchop
added
enhancement
kocchop
force-pushed
the
faysal/add-thd-2-ring-attn
branch
2 times, most recently
from
2935ff5 to
132bf4c
Compare
Enable CP + packing for context_parallel_strategy="ring" with load
balancing. On GPU, uses Transformer Engine's striped reorder for
THD-packed sequences. On TPU/CPU, falls back to pure-JAX reorder_sequence
and never imports TE.
Changes:
- common_types: Add ReorderStrategy enum (AUTO, DUAL_CHUNK_SWAP, STRIPED).
- configs: Add context_parallel_reorder_strategy (default "auto"). Reject
explicit STRIPED on non-GPU at config validation time.
- attention_op: Thread segment_positions through apply_attention,
cudnn_flash_attention, and __call__. Use segment_positions in TE's
SequenceDescriptor for packing. Restrict packing+CP to load-balanced
ring only. Note TE version constraint.
- attentions.py, attention_mla.py, gpt3.py: Pass inputs_positions into
attention_op calls (None for gpt3).
- max_utils: Hardware-dispatched reorder_causal_load_balanced. GPU uses
TE's reorder_causal_load_balancing; TPU/CPU uses reorder_sequence.
TE import is lazy and GPU-only.
- maxtext_utils: Thread reorder_strategy and hardware through
shard_reorder_causal_load_balanced and get_reorder_callable. Default
hardware="tpu" never triggers TE import.
- train_utils: Allow ring+packing; forbid all_gather+packing and
synthetic+packing. Resolve AUTO->STRIPED for packing else
DUAL_CHUNK_SWAP. Pass config.hardware to reorder callable. Build
data_loader after reorder wrapper is applied.
- attention_test_util: Pass cfg_cp.hardware so TPU tests use pure-JAX
reorder. Helper is TPU-oriented and does not model GPU packed behavior.
- tests: Add test_gpu_ring_attention_with_packing (sm90+).
Requires TE with reorder_causal_load_balancing; works with TE <=2.11 or
>=2.14 (incompatible with 2.12 and 2.13 due to a known bug).
kocchop
force-pushed
the
faysal/add-thd-2-ring-attn
branch
from
132bf4c to
eb52c0b
Compare
copybara-service
Bot
merged commit 9608068
into
AI-Hypercomputer:main
22 of 23 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Enables sequence packing for context parallelism with
ringstrategy using TransformerEngine's DotProductAttention. Includes comprehensive GPU tests for ring attention with packing for sm90+.reorder_causal_load_balancingapiCurrent support matrix for context parallelism on GPU:
all_gatherdual_chunk_swapringdual_chunk_swap(non-packed)striped(packed)Tests
Added a GPU integration test that works for sm90+.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.