[PyTorch][Attention] THD P2P context-parallel regression when padded cu_seqlens are value-equal but not object-identical

4cd705b75394563c0246bdddfa5d3148106c9285 のコミットで導入された PyTorch DotProductAttention の padding 検出器から着手し、get_cu_seqlens_on_cp_rank と THD の dQ/dK/dV の末尾をゼロ化するパスを調査します。提供されたベンチマークを使って pad_between_seqs を交互に指定するケースを再現し、カーネルの処理量と実行時間を比較します。値が同等なメタデータによって不要な padding 処理を回避しつつ、数値的な動作とグラフセーフな実行が維持されれば完了です。

Describe the bug

Commit 4cd705b75394563c0246bdddfa5d3148106c9285 introduces a PyTorch DotProductAttention performance regression under this combination:

• qkv_format="thd"
• P2P context parallelism
• pad_between_seqs=None (automatic detection)
• both regular and padded cumulative sequence-length tensors are provided
• the padded and unpadded tensors are distinct objects but have identical relevant values, so there is no actual padding between sequences

This is observable during ordinary eager training; CUDA graph capture does not need to be enabled.

The new automatic detection uses tensor object identity as a proxy for padding semantics:

if cu_seqlens_q_padded is cu_seqlens_q:
    pad_between_seqs = False
elif cu_seqlens_q_padded is not None or cu_seqlens_kv_padded is not None:
    pad_between_seqs = True

Thus independently allocated but value-identical tensors are classified as pad_between_seqs=True. With P2P context parallelism, that selects the path that repeatedly calls get_cu_seqlens_on_cp_rank, rather than the cheaper no-inter-sequence-padding path.

The same commit also adds THD dQ/dK/dV tail-zeroing operations. They launch arange/compare/masked-fill work even when the valid endpoint already equals the tensor endpoint and there is no tail to clear.

Steps/Code to reproduce bug

1. Create a BF16 DotProductAttention module with qkv_format="thd" and a four-rank P2P context-parallel group.

2. Provide independently allocated cumulative sequence-length tensors with identical values:

   cu_seqlens = torch.tensor([0, sequence_length], dtype=torch.int32, device="cuda")
   cu_seqlens_padded = cu_seqlens.clone()
   
   assert cu_seqlens_padded is not cu_seqlens
   assert torch.equal(cu_seqlens_padded, cu_seqlens)
   

3. Run repeated attention forward/backward calls, alternating these two cases in the same process:

   • automatic detection: pad_between_seqs=None
   • known-correct metadata: pad_between_seqs=False

4. Discard warmup and compare steady-state timings. A single attention forward/backward call shows a small direct overhead. The impact becomes much larger in an attention-heavy training schedule where the branch is exercised repeatedly and interacts with context-parallel stream scheduling.

We also performed a controlled source-level reverse experiment on Transformer Engine 2.18.0+27486e03. All arms used the same process, allocation, inputs, configuration, and byte-identical compiled Transformer Engine extensions; only the Python attention hunks from the cited commit differed. Each arm used 50 post-warmup iterations.

Variant	Mean iteration time	Median	Delta vs. stock
Stock	595.084 ms	589.000 ms	—
Revert padding detector only	560.958 ms	556.800 ms	-5.735%
Revert gradient zero-fill only	577.764 ms	574.950 ms	-2.911%
Revert both	554.254 ms	549.650 ms	-6.861%

An ABBA repetition of stock and the full reverse patch measured a 6.276% aggregate iteration-time improvement with the reverse patch. The stock and reverse-patched order drift was 1.534% and 0.861%, respectively.

Across all ranks in two Nsight Systems trials, the full reverse patch reduced the five-step trace span by 7.411% on average, with every paired rank faster. Over five steps it removed, per rank:

• 2,400 helper-generated kernel launches associated with get_cu_seqlens_on_cp_rank
• 900 masked-fill kernels from the new gradient tail-zeroing blocks

The source reversal reduced main-stream kernel work by 31.702 ms/rank and main-stream gaps by 237.154 ms/rank over the captured five-step window. These are separate trace observations, not additive wall-time attribution. Numerical-health checks remained clean.

Expected behavior

When the padded and unpadded cumulative sequence-length tensors have equal relevant values, automatic detection should not select the inter-sequence-padding path solely because they are different Python objects.

Could the API carry graph-safe padding metadata explicitly, or otherwise avoid using object identity as the semantic proxy? The tail-zeroing work could also be gated when metadata establishes that no gradient tail exists, while preserving CUDA graph compatibility.

The current workaround for callers that know there is no inter-sequence padding is to pass pad_between_seqs=False explicitly.

Environment overview

• Environment location: containerized bare-metal system
• Transformer Engine: 2.18.0+27486e03
• Installation: preinstalled container package

Environment details

• Python: 3.12.3
• PyTorch: 2.13.0a0+8145d630e8.nv26.6.54250401
• CUDA reported by PyTorch: 13.3
• cuDNN: compiled against 9.23; node-visible runtime 9.21.1

The causal comparison used one unchanged environment, so the cuDNN packaging detail was identical across all variants.

Device details

• 4x NVIDIA H100 80GB HBM3

Additional context

The detector is the primary contributor. After removing the zero-fill blocks, reverting the detector still improved iteration time by 4.069%. Once the detector was corrected, removing zero-fill added another 1.195%. The effects overlap on the same context-parallel critical path and therefore should not be added independently.