[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, backends.py, when "padding" is in the mask type and the user passed cumulative sequence lengths rather than a mask).

It was written as a loop over the batch that used 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 a design decision that is being second-guessed here. The construction was written as reference-mask code on the test side in NVIDIA#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 NVIDIA#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 was checked exhaustively against the previous implementation 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 produces 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: those suites pass unchanged.

Split out of a larger torch.compile change for DotProductAttention, where the same loop is also untraceable by dynamo (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

🤖 Generated with Claude Code

[pggPL]

Superseded by NVIDIA#3268, opened upstream with the same change squashed into a single commit.