THD + CP (p2p): pad_between_seqs auto-detect is blind to tail padding → nondeterministic forward, silently wrong attention (dropped real keys, exact-zero real rows), nondeterministic gradients

[JackRao123]

Summary

With qkv_format="thd", attn_mask_type="padding_causal", context parallelism (cp_comm_type="p2p"), and a sequence whose real length is not divisible by 2*cp_size (tail-padded — e.g. real 698 padded to 704 at CP4, the common case for packed THD training), the pad_between_seqs auto-detect in DotProductAttention.forward compares cu_seqlens_padded[:-1] with cu_seqlens[:-1] — deliberately ignoring tail padding after the last sequence (dot_product_attention.py @ v2.16, L1519–1529).

The call therefore takes the pad_between_seqs=False fast path in the CP p2p orchestration, which approximates per-ring-step seqlens as cu_seqlens // cp_size (context_parallel.py @ v2.16, L762–773, cp_p2p_fwd_prepare_qkv). That uniform-average approximation mislabels each rank's chunk-boundary rows whenever real % (2*cp) != 0. In short: the detect conflates "no padding between sequences" with "no padding at all" — an equivalence that only holds when CP is off, because CP wraps tail padding into a rank's local buffer and divides boundaries by cp_size.

Affects: every released TE we checked (2.16.0, 2.17, 2.17.1). main is only incidentally and partially unaffected — see below.

Reproduced on: TE 2.16.0, torch 2.11.0+cu130, cuDNN 9.19.0, sm103 (B300), 4 GPUs. cuDNN-independent: measured 18-cell matrix — TE {2.16.0, 2.17.1} fire 3/3 seeds × 60 iters on every cuDNN in {9.19.0.56, 9.21.1.3, 9.24.0.43, 9.25.0.15}, with the wobble at the identical row in 108/108 firing instances; TE main (2.19.0.dev0+8260f49) quiet 0/3 on all four (loaded libcudnn verified per cell via /proc/self/maps + cudnnGetVersion). Fires at CP4 within 20 iterations; CP2 more rarely.

Consequences (all proven at element level, fixed inputs)

1. Nondeterministic forward output at exactly one row per affected rank (local row T_LOCAL-2). The per-step softmax-LSE aux tensor is allocated with at::empty (allocateSpace(..., init_to_zeros=false), csrc/extensions/attention.cpp), and the cuDNN kernel writes only rows the (wrong) cu_seqlens calls real. The correction kernels (thd_second_half_lse_correction, thd_out_correction) iterate PADDED ranges and log-sum-exp-merge the uninitialized rows. The design contract — garbage LSE is harmless because out_per_step == 0 at padding rows (explicit zero-guard in thd_out_correction_kernel) — breaks at the one row that is mislabeled-padding for the diagonal/lower steps but REAL for the upper-triangle steps: a real output row gets scaled by exp(lse_real − merge(garbage, lse_real)).
2. Silently wrong attention (deterministic). The mislabeled cu also drops each rank's last two chunk rows as padding KEYS, so 2*(cp-1) real keys are invisible to every query (measured: 607/698 real rows differ from the CP1 reference by >1e-2, mean |Δ| 0.04 on random bf16 data), and real QUERY rows at the boundary get zero or partial outputs (cp-1 real tokens receive EXACT-ZERO attention output; more receive upper-steps-only contributions). Because this part is deterministic, it is identical on both sides of save-vs-reload determinism compares and cancels — it ships silently in every CP>1 THD training run with real % (2*cp) != 0.
3. Nondeterministic gradients on ALL CP ranks. The merged (garbage-bearing) LSE is saved for backward; the dKV ring spreads the contamination, so backward is nondeterministic across many rows on every rank — including ranks whose forward looks clean.

Minimal repro

Pure te.DotProductAttention — no Megatron, no model weights. THD + padding_causal + CP p2p + MQA (8q/1kv/d128, bf16), one 698-token sequence padded to 704, identical inputs each iteration, 4 GPUs, ~2 min:

torchrun --standalone --nproc-per-node=4 standalone_te_cp4_repro.py

Observed: CP ranks 0–2 produce 3–7 distinct bitwise outputs in 20–30 iterations, always differing at exactly local row 174 (T_LOCAL-2, the tail of the rank's second load-balanced chunk); rank cp-1 clean in forward (it has no upper-triangle steps). A --backward variant shows grads nondeterministic on 4/4 ranks. CP1 is deterministic.

standalone_te_cp4_repro.py

#!/usr/bin/env python3
"""Standalone repro: TE cuDNN FusedAttention forward nondeterminism
under context-parallel p2p, THD layout. No megatron, no model weights.

  qkv_format=thd, per-rank q heads=8, kv heads=1 (MQA slice), head_dim=128,
  attn_mask_type=padding_causal, window=(-1,0), bf16, cp_size=4 (p2p ring),
  one packed sequence of 698 tokens padded to 704 (=2*cp*88), softmax scale
  default, no dropout, forward only (inference_mode).

Each CP rank holds load-balanced chunks [r, 2*cp-1-r] of 88 tokens each.
Runs N forwards on identical inputs and compares outputs bitwise across
iterations; reports distinct-count and differing rows per rank.

Launch (on one 8-GPU node, uses 4):
  torchrun --standalone --nproc-per-node=4 standalone_te_cp4_repro.py
"""
from __future__ import annotations

import hashlib
import os
import sys

import torch
import torch.distributed as dist

import transformer_engine.pytorch as te


SEQ_REAL = 698
T_TOTAL = 704       # padded; must be divisible by 2*CP
HEADS_Q = 8
HEADS_KV = 1
HEAD_DIM = 128
N_ITERS = 20
SEEDS = [16, 17, 18]


def sha(t: torch.Tensor) -> str:
    return hashlib.sha256(
        t.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes()
    ).hexdigest()[:16]


def main() -> int:
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    world = dist.get_world_size()
    CP = world
    CHUNK = T_TOTAL // (2 * CP)
    T_LOCAL = 2 * CHUNK
    torch.cuda.set_device(rank)
    dev = torch.device("cuda", rank)
    if rank == 0:
        import transformer_engine
        print(f"CP={CP} chunk={CHUNK} TE={transformer_engine.__version__} "
              f"torch={torch.__version__} cudnn={torch.backends.cudnn.version()} "
              f"NVTE_ALLOW_NONDETERMINISTIC_ALGO={os.environ.get('NVTE_ALLOW_NONDETERMINISTIC_ALGO')}",
              flush=True)

    cp_group = dist.new_group(ranks=list(range(CP)), backend="nccl")
    cp_stream = torch.cuda.Stream(device=dev)

    attn = te.DotProductAttention(
        num_attention_heads=HEADS_Q,
        kv_channels=HEAD_DIM,
        num_gqa_groups=HEADS_KV,
        attention_dropout=0.0,
        attn_mask_type="padding_causal",
        qkv_format="thd",
        softmax_scale=None,
    ).to(dev)
    attn.set_context_parallel_group(
        cp_group, list(range(CP)), cp_stream, cp_comm_type="p2p"
    )

    t_total = 2 * CP * CHUNK  # 704
    cu_q = torch.tensor([0, SEQ_REAL], dtype=torch.int32, device=dev)
    cu_q_padded = torch.tensor([0, t_total], dtype=torch.int32, device=dev)

    overall_rc = 0
    for seed in SEEDS:
        g = torch.Generator(device="cpu").manual_seed(seed)
        qf = torch.randn(t_total, HEADS_Q, HEAD_DIM, generator=g).bfloat16()
        kf = torch.randn(t_total, HEADS_KV, HEAD_DIM, generator=g).bfloat16()
        vf = torch.randn(t_total, HEADS_KV, HEAD_DIM, generator=g).bfloat16()

        # load-balanced CP shard: chunks [rank, 2*CP-1-rank]
        idx = torch.cat(
            [
                torch.arange(rank * CHUNK, (rank + 1) * CHUNK),
                torch.arange((2 * CP - 1 - rank) * CHUNK, (2 * CP - rank) * CHUNK),
            ]
        )
        q = qf[idx].to(dev).requires_grad_(False)
        k = kf[idx].to(dev)
        v = vf[idx].to(dev)

        outs = []
        with torch.inference_mode():
            for _ in range(N_ITERS):
                o = attn(
                    q, k, v,
                    cu_seqlens_q=cu_q,
                    cu_seqlens_kv=cu_q,
                    cu_seqlens_q_padded=cu_q_padded,
                    cu_seqlens_kv_padded=cu_q_padded,
                    max_seqlen_q=t_total,
                    max_seqlen_kv=t_total,
                )
                outs.append(o.clone())

        hashes = [sha(o) for o in outs]
        distinct = len(set(hashes))
        rows = set()
        if distinct > 1:
            base = outs[0].view(T_LOCAL, -1).float()
            for o in outs[1:]:
                d = (o.view(T_LOCAL, -1).float() - base).abs().amax(dim=1)
                rows.update(torch.nonzero(d > 0).flatten().tolist())
        print(
            f"[seed {seed}] rank {rank}: {distinct} distinct outputs over "
            f"{N_ITERS} iters; wobble rows: {sorted(rows)[:10]}",
            flush=True,
        )
        flag = torch.tensor([1 if distinct > 1 else 0], device=dev)
        dist.all_reduce(flag)
        if rank == 0:
            verdict = "NONDETERMINISTIC" if flag.item() > 0 else "deterministic"
            print(f"[seed {seed}] GLOBAL: {verdict} ({flag.item()}/{CP} ranks)",
                  flush=True)
        if flag.item() > 0:
            overall_rc = 1

    dist.destroy_process_group()
    return overall_rc


if __name__ == "__main__":
    sys.exit(main())

Evidence chain (all element-level, fixed inputs)

• Per-ring-step fingerprints: every q/k/v/cu INPUT and every kernel OUT is bitwise stable across iterations; ONLY the per-step LSE varies, at exactly the mislabeled rows (diagonal step rows {174,175}; upper-triangle step row {87}), then total-LSE columns {174,175}, then out row {174} via the upper-step thd_out_correction calls only. cuDNN's written values are fully deterministic — this is not a cuDNN bug.
• Causality (poison test): overwriting the step-LSE padding rows with a constant makes the output DETERMINISTIC: +1000 → deterministic and wrong (row 174 becomes exactly 0, as the scale factor underflows); -1e30 → deterministic (but still missing the never-computed diagonal-step contribution at that row). Stock wobble + both poisons pinning the output ⇒ uninitialized LSE content is THE cause.
• Fix validation: passing pad_between_seqs=True (engaging the exact get_cu_seqlens_on_cp_rank path) makes forward AND backward bitwise deterministic (100/100 iters × 3 seeds × CP{2,4}) and restores correctness: reassembled CP4 output matches CP1 within bf16 rounding (max |Δ| 0.0039, no row >1e-2).
• Version matrix (18/18 cells, per-cell loaded-cuDNN proof): TE 2.16.0 and 2.17.1 fire on ALL of cuDNN 9.19/9.21.1/9.24/9.25; TE main quiet on all four — cuDNN is not the variable; the main auto-detect rewrite is the difference. NVTE_FUSED_ATTN_DIRECT_SEQLENS arms inert (symbol absent in all three builds including main).
• Trainer-scale validation (Megatron-based stack, dense 0.6B model, TP2×CP4): stock nondeterministic across 10 forwards (max |Δ| 0.083); with pad_between_seqs=True bitwise-deterministic 10/10. On a 550B hybrid MoE production model the same defect amplified through router top-k flips to |Δ| up to 5 — not a cosmetic wobble at scale.

Why main is only incidentally (and partially) unaffected

main rewrote the auto-detect in #2898 (merged 2026-07-01; motivation: torch.equal syncs during CUDA-graph capture — not a correctness fix for this bug). The new detect is an identity check: distinct padded/real cu tensor objects ⇒ pad_between_seqs=True. Callers passing distinct tensors therefore land on the exact path on main.

But main's own comment codifies the blind spot as intended behavior (dot_product_attention.py @ main, L1950–1969):

"If padded cu_seqlens are the same object as the unpadded ones, no real inter-sequence padding exists (only THD tail padding) -- treat as False."

Under CP p2p that inference is exactly wrong: tail-padding-only is precisely the case that breaks cu_seqlens // cp_size whenever real % (2*cp_size) != 0. Callers that pass the SAME tensor object for real and padded cu — a pattern used in the wild (e.g. verl, see #2892) — still reach the broken fast path on main today, with all three defects above. The // cp_size fast path itself is unchanged on main (context_parallel.py @ main, L721), and the LSE aux tensor is still allocated uninitialized (attention.cpp @ main, L279–283).

Asks

1. Under CP, tail padding must count: make the detect/route CP-aware (e.g. compare FULL cu arrays where value-compares are acceptable, or route THD+CP+padded-cu to the exact path), or document that THD+CP requires real % (2*cp) == 0 and assert.
2. Add a regression test pinning these semantics (CP{2,4}, tail-padded THD: bitwise-determinism across iterations + equality with CP1 reference), so the incidental main fix can't regress.
3. Consider zero/neg-inf-initializing the LSE aux tensor (or masking padded rows post-kernel) as defense in depth — the zero-guard contract in thd_out_correction_kernel is fragile. (Note this converts the nondeterminism to deterministic-wrong; it is NOT a correctness fix by itself.)

Prior reports of this symptom (unanswered) and related items

• How to pad sequences that cannot be evenly split in context parallelism (CP)? #1929 and Megatron-LM #1669 (same reporter, 2025): seq len 7 at CP2 padded to 8; TE "internally behaves as if cu_seqlens=[0,6]"; the repro shows a REAL query row with exact-zero CP output — our defect (2). Filed as a usage question, never answered; the mechanism above is the answer.
• [BUG] SIGSEGV in fused_attn_fwd when cu_seqlens_q != cu_seqlens_q_padded with FP8 blockwise #2892: documents the same-object aliasing pattern (cu_seqlens passed as cu_seqlens_padded) in verl — i.e. main's identity check does not cover real callers.
• [Bugfix][PyTorch][FlashAttention][CP] Honor requested FA padding in CP tests #3269: precedent that a wrong pad flag leaves CP padding uninitialized ("Honor requested FA padding in CP tests").
• [PyTorch] Add pad_between_seqs support for non-CP and CP (A2A and P2P) with FA3 + THD (varlen) #2596: the exact-path infrastructure (pad_between_seqs support for CP A2A+P2P) that the fix routes to.
• Fused attention with THD format + CP may cause a bug in backward pass, leading to NaN values during training LLM. #2186: THD+CP second-chunk-tail NaN in backward, fixed in cuDNN 9.18 — different mechanism (cuDNN backward workspace init), same structural site.
• Inconsistent fused_attn_bwd results of dk, dv under THD layout + GQA on multiple calls #2448: adjacent evidence that uninitialized/reused aux memory in THD attention produces silent nondeterministic wrongness (CP=1, different mechanism).

Workaround for users on released TE

Pass pad_between_seqs=True explicitly for THD+CP calls with tail padding (public kwarg, available in all affected releases). Cost measured at ~0.3–1.5 ms/call host-side at microbench scale; noise at training scale.

[JackRao123]

Closing: addressed on our side by passing pad_between_seqs=True explicitly for THD+CP with tail padding. Mechanism, repro, and workaround above remain valid for anyone hitting this on released TE (<=2.17.1).