[PyTorch] Regression test + docs for THD CP tail-padding mislabeling (pad_between_seqs)

[JackRao123]

Description

Regression test + documentation for the THD + context-parallel tail-padding defect reported in #3331.

In short: with qkv_format="thd" + CP p2p, a sequence whose real length is not divisible by 2*cp_size (tail-padded — the common case for packed THD training) is mishandled whenever the pad_between_seqs=False fast path is taken, because that path approximates per-ring-step seqlens as cu_seqlens // cp_size and thereby mislabels chunk-boundary rows. Consequences (proven at element level in #3331): nondeterministic forward rows, silently wrong attention (real keys dropped, real query rows exact-zero), nondeterministic gradients on all CP ranks.

main routes callers that pass distinct padded/unpadded cu_seqlens tensors to the exact path (incidental effect of the #2898 auto-detect rewrite), but nothing pins that behavior, and the docstring/comment still describe tail padding as a safe False signal — which is incorrect under CP.

Refs #3331 (this PR is the regression test + docs; the deeper routing question — e.g. whether THD+CP should always take the exact per-rank path, and LSE aux hardening as defense in depth — is left for discussion on the issue, per the asks there).

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue) — regression coverage for a released bug

Changes

• tests/pytorch/attention/run_cp_thd_tail_padding.py (new): distributed runner. Single packed sequence, real length 698 (not divisible by 2*cp for cp∈{2,4}), tail-padded; FusedAttention backend, padding_causal, THD, CP p2p. Three arms per pad mode (pad_between_seqs=True explicit, and auto-detect with distinct padded tensors): (1) forward bitwise determinism over 25 iterations, (2) fwd+bwd bitwise determinism over 10 iterations, (3) correctness of out/dq/dk/dv vs an unpadded no-CP reference on all real rows.
• tests/pytorch/attention/test_cp_thd_tail_padding.py (new): pytest wrapper launching the runner at world sizes 2 and 4 via run_distributed.
• dot_product_attention.py: docstring for pad_between_seqs and the auto-detect comment now state that tail padding is not benign under CP, and that tail-padded THD+CP callers must pass pad_between_seqs=True explicitly or supply distinct padded cu_seqlens tensors.

Validation (H100×4, torch 2.11.0+cu128, cuDNN 9.25.0.15)

run	TE 2.17.1 (affected release)	TE main (this branch)
CP4	FAIL — auto arm: forward nondeterministic 24/24 iters; fwd+bwd nondeterministic; out/dq/dk/dv vs reference max |Δ| up to 1.41	PASS (both modes)
CP2	FAIL — auto arm: deterministic mis-attention vs reference, max |Δ| up to 0.95	PASS (both modes)

The explicit pad_between_seqs=True arm passes on 2.17.1 as well, confirming the documented workaround.

One environment note: TE 2.17.1's THD backward fails wholesale with cuDNN 9.19.0 (the version torch 2.11.0+cu128 pins) — CUDNN_STATUS_BAD_PARAM at fused_attn_f16_arbitrary_seqlen.cu:934 (reshape attribute), independent of this defect and reproducible on a plain no-CP THD backward; cuDNN 9.25.0.15 resolves it. May be worth a minimum-cuDNN note for the THD backward path.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes — new test validated on H100×4 (see table); it exercises a path no existing test covers (existing THD CP tests pad every sequence, so internal boundaries always move and the auto-detect never sees the tail-only case)

[greptile-apps]

Greptile Summary

The PR documents the THD context-parallel tail-padding constraint and adds a distributed numerical regression test for explicit and auto-detected padding modes.

• Adds deterministic forward/backward and no-CP reference comparisons at context-parallel sizes 2 and 4.
• Clarifies that tail padding requires the exact per-rank path under context parallelism.
• The new regression test still needs to be wired into a repository QA job.

Confidence Score: 4/5

The PR appears safe to merge after the non-blocking CI coverage omission is addressed.

The test and documentation changes are internally consistent, but the regression test is not selected by any checked repository QA invocation and therefore does not currently protect the behavior in CI.

Files Needing Attention: tests/pytorch/attention/test_cp_thd_tail_padding.py and the relevant QA test script

Important Files Changed

Filename	Overview
tests/pytorch/attention/run_cp_thd_tail_padding.py	Adds a focused distributed runner validating deterministic outputs and gradients against an unpadded no-CP reference.
tests/pytorch/attention/test_cp_thd_tail_padding.py	Adds CP2/CP4 pytest entry points, but the file is absent from the repository's explicit CI test lists.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Updates documentation and comments to explain why tail padding is unsafe on the approximate context-parallel path.

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[greptile-apps]

Regression test omitted from CI

The repository's QA jobs invoke attention tests through explicit file paths, but none includes test_cp_thd_tail_padding.py. The new regression therefore does not run in CI, allowing the THD context-parallel tail-padding defect to be reintroduced without failing a job; please add this file to the relevant distributed QA invocation.

Knowledge Base Used: Tests and QA

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[JackRao123]

Closing for now — we'll carry the workaround on our side (explicit pad_between_seqs=True). The analysis in #3331 stands; happy to reopen if there's maintainer interest in the regression test.