[PyTorch] Build the attention padding mask without host syncs

[pggPL]

Description

get_padding_mask turns cu_seqlens into the boolean padding mask that UnfusedDotProductAttention consumes (its only caller, in backends.py, when "padding" is in the mask type and the caller supplied cumulative sequence lengths rather than a mask).

It is written as a loop over the batch that uses each sequence length as a Python list multiplier:

for i in range(batch_size):
    attention_mask_q = torch.cat([
        attention_mask_q,
        torch.Tensor([False] * seqlens_q[i] + [True] * (max_seqlen_q - seqlens_q[i]))...
    ], dim=0)

seqlens_q[i] is an element of a CUDA tensor, so [False] * seqlens_q[i] forces an implicit .item(): one device synchronization per sequence, plus one torch.cat per sequence, on every forward pass that takes this path.

This is not second-guessing a design decision. The construction was written as reference-mask code on the test side in #818, where a host round-trip per sequence costs nothing, and was lifted into the library when get_padding_mask was introduced for KV caching in #1355. Its twin is still in tests/pytorch/attention/test_attention.py and stays there — on the test side it is perfectly fine. This PR only removes the consequence of that move onto a per-forward path.

The mask is positions >= seqlens, which needs no host round-trip:

seqlens = cu_seqlens[1 : batch_size + 1] - cu_seqlens[:batch_size]
positions = torch.arange(max_seqlen, device=cu_seqlens.device)
return (positions.unsqueeze(0) >= seqlens.unsqueeze(1)).view(batch_size, 1, 1, max_seqlen)

Type of change

• Bug fix (non-breaking change which fixes an issue)
• Code refactoring

Changes

• get_padding_mask: vectorized, device-side construction; no synchronization, one kernel instead of batch_size iterations.
• The slice [1 : batch_size + 1] is deliberate and preserves existing behaviour: cu_seqlens may hold more than batch_size + 1 entries, because inference allocates it for the maximum batch size. The old loop ignored the extra entries by iterating range(batch_size).
• The mask now lands on cu_seqlens.device instead of a hard-coded "cuda". DotProductAttention asserts CUDA inputs, so this is not reachable from there.

Equivalence

The input space is discrete and small, so equivalence with the previous implementation was checked exhaustively rather than sampled — values, shape, dtype and device:

• self attention: every combination of sequence lengths in [0, max_seqlen] for max_seqlen 1..8 and batch_size 1..4, with and without a cu_seqlens buffer longer than batch_size + 1 — 35368 cases, 0 mismatches;
• cross attention: every combination of q and kv lengths for max_seqlen_q, max_seqlen_kv 1..5 and batch_size 1..3 — 202100 cases, 0 mismatches.

The only inputs where the two differ are invalid ones: a sequence longer than max_seqlen made the old code raise from torch.cat (rows of unequal length), and now yields an all-False row.

Testing

No new tests. The function is already exercised by the thd and padding cases in tests/pytorch/attention/test_attention.py (the -k thd subset alone calls it 10 times, with numerics compared against a reference) and by tests/pytorch/attention/test_kv_cache.py, which is what covers the longer-than-batch cu_seqlens buffers. Verified on RTX Ada (sm89): test_attention.py, test_kv_cache.py and test_gqa.py pass with no new failures relative to main.

This is split out of a larger torch.compile change for DotProductAttention, which will be submitted separately — there the same loop is additionally untraceable by dynamo, because the sequence lengths become data-dependent (unbacked) SymInts. This part stands on its own as an eager-mode fix and is reviewable without any torch.compile context.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works — covered by existing tests, see Testing
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

The PR vectorizes padding-mask construction to avoid per-sequence host synchronization while preserving existing batch slicing and CUDA mask placement.

• Computes sequence lengths and padding positions with device-side tensor operations.
• Retains support for cumulative-length buffers larger than the active batch.
• Transfers the completed mask to the current CUDA device for use by unfused attention.

Confidence Score: 5/5

The PR appears safe to merge.

The previously reported device mismatch is resolved by transferring the completed padding mask to CUDA before it is consumed, and no blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Replaces iterative host-synchronizing mask construction with vectorized tensor operations and resolves the previously reported CPU-mask device mismatch.

_{Reviews (2): Last reviewed commit: "Keep the padding mask on the attention d..." | Re-trigger Greptile}

[pggPL]

/te-ci pytorch