[PyTorch] Support padded THD sequences with FlashAttention 4

[xiaoyao0115]

Description

Add support for pad_between_seqs=True with the FlashAttention 4 backend for physically padded THD tensors.

The FlashAttention varlen kernels use padded cumulative sequence boundaries to describe the physical tensor layout and seqused_q/seqused_k to describe the valid token counts. FlashAttention 4 was previously disabled for this mode, and FlashAttention 3/4 may leave unused padding rows unwritten. This change enables the FA4 path and guarantees zero values at inter-sequence padding positions in both the attention output and Q/K/V gradients.

FA2 remains disabled for pad_between_seqs, and FA4 context parallel support is unchanged (not enabled by this PR).

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)

Changes

• Pass physical cu_seqlens_q/k_padded and valid seqused_q/k to the FA4 varlen API.
• Keep FA4 eligible during backend selection for THD tensors with padding between sequences.
• Mask non-CP FA3/FA4 output padding rows and mask padding rows in Q/K/V gradients through an autograd identity.
• Extend DPA tests to require zero output and gradient values for every inter-sequence padding range.

Validation

• Direct FlashAttention 4 (4.0.0b4) kernel parity test on GB200 against an equivalent compact varlen reference, including output and Q/K/V gradients: passed.
• Output and Q/K/V gradient padding-zero assertions: passed.
• python -m py_compile on all modified Python files: passed.
• git diff --check: passed.

The available container's preinstalled Transformer Engine C++ extension does not match the current main Python API, so the full source-tree pytest suite could not be run in that environment. A separate FA3 runtime was not available; the shared masking path was exercised with FA4, while the existing FA3 pad_between_seqs path remains unchanged.

Checklist

• I have read and followed the contributing guidelines.
• The functionality in the stated scope is complete.
• I have commented code in hard-to-understand areas.
• My changes generate no new warnings in the checks run above.
• I have added tests that prove the change is effective.