[PyTorch][CP] Add THD format support for AllGather-based Context Parallelism#2829
Conversation
… cu_seqlens - Use per-step cu_seqlens_q_padded to select Q chunks instead of tensor slicing - Use padded cu_seqlens_kv for K/V reordering (ensures divisibility) - Add cu_seqlens_kv and cu_seqlens_kv_padded to AllGather function signature - Compute per-step Q and KV cu_seqlens correctly from actual seqlens - Support non-causal attention (all KV visible) - Zero-initialize out/dq for THD to avoid garbage in padding regions - Save per-step cu_seqlens in ctx for backward (avoid recomputation) Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Remove skip gates that blocked THD format with all_gather CP comm type. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
…seqlens_q_padded The interleaved valid mask computation assumed cu_seqlens_q_padded starts at 0. With the CP offset-based approach, cu_seqlens_q_padded can start at a non-zero offset, causing a size mismatch. Use absolute positions from cu_seqlens_q_padded to build the valid mask instead. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
sudhakarsingh27
force-pushed
the
cp_thd_swa_with_ag
branch
from
1164a15 to
b4db9eb
Compare
for more information, see https://pre-commit.ci
| if qkv_format == "thd": | ||
| # [cp*t, h, d] -> reorder to contiguous per-sequence order -> [t_full, h, d] | ||
| chunk_ids_for_kv_ag = get_seq_chunk_ids_for_reordering_before_attn(cp_size, k.device) | ||
| k_ag = reorder_seq_chunks_after_a2a_before_attn_thd( |
There was a problem hiding this comment.
This reorder_seq_chunks_after_a2a_before_attn_thd and the other releated method are not "a2a" specific now, rename them to something like dualchunk_to_contiguous_order_thd and the other one contiguous_to_dualchunk_order_thd
There was a problem hiding this comment.
Resolved on the current branch, with final cleanup in 0e926c42. The THD reorder entry points are now reorder_thd_sequences_to_rank_sharded and reorder_thd_sequences_to_contiguous, and the stale Python permutation helpers were removed. Both wrappers call the fused tex.thd_reorder path, so this logic is no longer A2A-named or A2A-specific.
There was a problem hiding this comment.
Still resolved at PR head af2bd1c3. The Python wrappers are not A2A-specific (reorder_thd_sequences_to_contiguous / reorder_thd_sequences_to_rank_sharded) and now call the renamed fused binding tex.thd_cp_reorder_sequences.
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
…formerEngine into cp_thd_swa_with_ag
for more information, see https://pre-commit.ci
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
for more information, see https://pre-commit.ci
sudhakarsingh27
changed the title
…formerEngine into cp_thd_swa_with_ag
Greptile SummaryThis PR adds THD (variable-length sequence) format support to
Confidence Score: 4/5Safe to merge for environments with FusedAttention or FlashAttention v3; a user on FA2-only hardware who reaches the AllGather+THD forward path receives silently wrong attention output with no error. The AllGather+THD forward and backward use per-step cu_seqlens_q_padded offsets that FA2 varlen cannot represent. Both steps read from Q position 0, producing identical partial outputs. The test file skips this combination with pytest.skip, but production code has no corresponding assert — so any direct caller without FA3 or cuDNN reaches the broken path silently. All other changes (stream sync, max-logit masking, CUDA helper kernels, pybind wrappers) look correct and well-tested. transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py — specifically the not use_fused_attention + not use_flash_attn_3 + qkv_format == 'thd' path in AttnFuncWithCPAndKVAllGather.forward (and the symmetric backward path). Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Fwd as AttnFuncWithCPAndKVAllGather.forward
participant AG as AllGather (cp_group)
participant Reorder as thd_cp_rank_order_to_sequence_order
participant Stream0 as current_stream (step 0)
participant Stream1 as cp_stream (step 1)
participant Copy as thd_copy_valid_tokens
Fwd->>AG: gather_along_first_dim(k, v)
AG-->>Fwd: "k_ag, v_ag [cp*t, h, d]"
Fwd->>Reorder: reorder k_ag, v_ag (THD path)
Reorder-->>Fwd: k_ag, v_ag in sequence order
Fwd->>Fwd: cp_stream.wait_stream(current_stream)
Fwd->>Stream0: step 0 attention (q_full, cu_seqlens_q_padded_step0)
Fwd->>Stream1: step 1 attention (q_full, cu_seqlens_q_padded_step1)
Stream0-->>Fwd: out_per_step[0]
Stream1-->>Fwd: out_per_step[1]
Fwd->>Copy: copy valid tokens from out_per_step[0] to out
Fwd->>Copy: copy valid tokens from out_per_step[1] to out
Copy-->>Fwd: out (accumulated rank-local result)
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant Fwd as AttnFuncWithCPAndKVAllGather.forward
participant AG as AllGather (cp_group)
participant Reorder as thd_cp_rank_order_to_sequence_order
participant Stream0 as current_stream (step 0)
participant Stream1 as cp_stream (step 1)
participant Copy as thd_copy_valid_tokens
Fwd->>AG: gather_along_first_dim(k, v)
AG-->>Fwd: "k_ag, v_ag [cp*t, h, d]"
Fwd->>Reorder: reorder k_ag, v_ag (THD path)
Reorder-->>Fwd: k_ag, v_ag in sequence order
Fwd->>Fwd: cp_stream.wait_stream(current_stream)
Fwd->>Stream0: step 0 attention (q_full, cu_seqlens_q_padded_step0)
Fwd->>Stream1: step 1 attention (q_full, cu_seqlens_q_padded_step1)
Stream0-->>Fwd: out_per_step[0]
Stream1-->>Fwd: out_per_step[1]
Fwd->>Copy: copy valid tokens from out_per_step[0] to out
Fwd->>Copy: copy valid tokens from out_per_step[1] to out
Copy-->>Fwd: out (accumulated rank-local result)
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
…ific helpers The AllGather THD path was not extending KV visibility beyond the causal boundary when window_size had a right component > 0, meaning tokens right of the diagonal were invisible to the kernel. Fix by adding window_size[1] to visible_padded (clamped at actual seqlen) and max_seqlen_kv_. Also rename reorder helpers to backend-neutral names since AllGather now uses them too, and add a clarifying comment for non-causal KV cu_seqlens. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
for more information, see https://pre-commit.ci
…formerEngine into cp_thd_swa_with_ag
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
| return out; | ||
| } | ||
|
|
||
| void thd_valid_copy(at::Tensor out, const at::Tensor &inp, const at::Tensor &cu_seqlens_padded, |
There was a problem hiding this comment.
Should cp_ be in the names as well? i.e. thd_cp_reorder_sequences/cp_thd_reorder_sequences or thd_cp_copy_valid_tokens/cp_thd_copy_valid_tokens? (Just a suggestion)
There was a problem hiding this comment.
Pushed in 27d2b84a, with pre-commit follow-up af2bd1c3: the Python-visible/C++ wrapper names are now thd_cp_reorder_sequences and thd_cp_copy_valid_tokens. The underlying C API names remain nvte_cp_thd_*.
| "Fused dual-chunk THD reorder for context parallel (gather/scatter), inline index", | ||
| py::call_guard<py::gil_scoped_release>()); | ||
| m.def("thd_valid_copy", &transformer_engine::pytorch::thd_valid_copy, | ||
| "Sync-free copy of valid THD token rows into an accumulator (CP AllGather fwd/bwd)", |
There was a problem hiding this comment.
What's "token rows" here? Sequences? Ranks?
Also, what's "inline index" in thd_reorder? :)
There was a problem hiding this comment.
Pushed in 27d2b84a, with pre-commit follow-up af2bd1c3: cleaned up the pybind docstrings. They no longer use “token rows” or “inline index”; they now describe reordering between CP rank-sharded dual-chunk order and contiguous per-sequence order, and copying valid THD sequence entries from a padded tensor.
|
I feel the code changes are complex enough that I have to rely on the tests. Please make sure the tests cover both Fused and Flash v3 backends, SWA/full attention/causal mask, pad_between_seqs=T/F combinations, and post the performance numbers if there's any (for the new reorder/valid_copy kernels vs the old PyTorch-based functions). Thanks! |
|
Also, please fix the DCO, and make the comments a bit more succinct if you can (some comments are quite thorough but also very long :) ). Thanks! |
Resolve review comments for PR 2829 by tightening the THD all_gather output shape, renaming the new THD CP helper bindings, removing the unrelated pybind helper extraction, and aligning the FP8 t3hd aux handling with the post-FP8DS code path. Also clean up the THD CP test skip logic and remove an unnecessary dtype conversion from the THD max-logit mask construction. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
for more information, see https://pre-commit.ci
The AllGather CP path already has a support assert that rejects padding masks for non-THD inputs. Keep the earlier THD-specific padding requirement, and rely on the later AllGather support assert for the non-THD padding case so the check is not duplicated. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
sudhakarsingh27
force-pushed
the
cp_thd_swa_with_ag
branch
from
af2bd1c to
46b764e
Compare
| ] | ||
| # Adjust chunks for each step | ||
| thd_cu_seqlens_kv_per_step[0][1:] = visible_actual[0].cumsum(0) | ||
| thd_cu_seqlens_kv_per_step[1][1:] = visible_actual[1].cumsum(0) |
There was a problem hiding this comment.
L3286-L3370 could be cudafied if this becomes a performance bottleneck.
|
/te-ci pytorch L3 |
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
…leanup Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
…codex/pr2829-review-cleanup Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
for more information, see https://pre-commit.ci
Rename the new THD CP reorder helpers so the source and destination layouts are encoded in the API names instead of a direction boolean. Also rename the valid-token copy helper to describe its per-split to rank-local accumulator role. Guard copy-valid tokens that precede the first padded THD offset before indexing shared cu_seqlens arrays; later split offsets can legitimately leave those token positions outside any valid sequence range. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> (cherry picked from commit 55f9f18747510d92791ef2650c65d93d9d90c27c)
…view-cleanup Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
for more information, see https://pre-commit.ci
|
/te-ci pytorch L3 |
Bring PR 2829 up to date with latest main (77054fa): picks up the EP common C API + NCCL EP backend, dense router output, VERSION bump to 2.18.0.dev0, and the multi_tensor_apply int32-overflow fix. Merge is conflict-free; the only files touched by both sides are header/binding decls that auto-merge. Local validation to follow. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
|
/te-ci pytorch L3 |
remove confusing info Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
3 participants
Description
Add THD (variable-length sequence) format support to
AttnFuncWithCPAndKVAllGather. Previously, AllGather-based CP only supported fixed-length formats (bshd/sbhd). THD format packs variable-length sequences into a single[t, h, d]tensor tracked bycu_seqlens, which is needed for workloads with heterogeneous sequence lengths.The key challenge is that AllGather CP splits Q across 2 steps (one per local chunk), but THD tensors cannot be naively sliced like fixed-length formats. This PR uses an offset-based approach: the full Q tensor is passed to the attention kernel each step, with per-step
cu_seqlens_q_paddedvalues directing the kernel to read the correct chunk. This avoids tensor slicing and follows the padded THD convention used by the backends.Type of change
Changes
cu_seqlens_q_paddedselects which chunk the kernel reads from the full Q tensor, instead of slicing Q per step.cu_seqlens_kv, and window ranges for causal, full/no-window, and SWA cases.pad_between_seqsandseqused_khandling.cu_seqlens_q_paddedin the valid-token mask without.item()D2H synchronizations.test_cp_utils.py.Checklist:
Latest validation update (2026-06-11)
CP THD AllGather coverage was validated across FusedAttention and FlashAttention v3.
False/TrueFalse/True(128,0)False/True(128,0)(512,512)test_cp_utils.pyalso passed 14/14 on both H100 and B200 L1 in CI pipeline 54405200. It covers the THD helper kernels, including reorder/copy-valid tests against the legacy Python reference paths.FusedAttention does not expose the FlashAttention-specific
pad_between_seqsaxis; Fused THD padding semantics are covered through the THD/padding mask path.Helper-kernel microbenchmarks, single H100, bf16,
batch=16,seqlen=4096,heads=16,dim=64:thd_cp_reorder_sequencescontiguous->rankthd_cp_reorder_sequencesrank->contiguousthd_cp_copy_valid_tokensNote: the helper-kernel numbers are fixed-shape microbenchmarks, not a sweep over sequence length or number of sequences.