From d050775067a04e1cb4147df5f5f6c52ae9b4ad4a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 28 Jul 2026 16:54:47 +0200 Subject: [PATCH 1/2] [PyTorch] Build the attention padding mask without host syncs get_padding_mask converts cu_seqlens into the boolean padding mask that UnfusedDotProductAttention consumes. It looped over the batch and used each sequence length as a Python list multiplier, which reads a GPU tensor on the host: one device synchronization per sequence, plus a torch.cat per sequence, on every forward pass that takes this path. Build the mask on the device instead: positions >= seqlens, broadcast over the batch. Equivalence with the previous implementation was checked exhaustively over every combination of sequence lengths for small batch sizes, for both self and cross attention. Signed-off-by: Pawel Gadzinski --- .../attention/dot_product_attention/utils.py | 62 +++++++------------ 1 file changed, 21 insertions(+), 41 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 6c47d0f1cb..3e728cfa4a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1652,51 +1652,31 @@ def get_padding_mask( max_seqlen_kv: int = None, attention_type: str = "self", ): - """Convert cu_seqlens to attention_mask""" + """Convert cu_seqlens to attention_mask. + + Built with device-side ops only: reading the sequence lengths on the host + would synchronize the device once per sequence. + """ assert ( cu_seqlens_q is not None and max_seqlen_q is not None ), "cu_seqlens_q and max_seqlen_q are required for self-attention and cross-attention" - seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] - attention_mask_q = torch.Tensor([]).to(dtype=torch.bool) - if attention_type == "cross": - assert ( - cu_seqlens_kv is not None and max_seqlen_kv is not None - ), "cu_seqlens_kv and max_seqlen_kv are required for cross-attention" - seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] - attention_mask_kv = torch.Tensor([]).to(dtype=torch.bool) - 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])) - .to(dtype=torch.bool) - .unsqueeze(0) - .unsqueeze(0) - .unsqueeze(0), - ], - dim=0, - ) - if attention_type == "cross": - attention_mask_kv = torch.cat( - [ - attention_mask_kv, - torch.Tensor([False] * seqlens_kv[i] + [True] * (max_seqlen_kv - seqlens_kv[i])) - .to(dtype=torch.bool) - .unsqueeze(0) - .unsqueeze(0) - .unsqueeze(0), - ], - dim=0, - ) - attention_mask_q = attention_mask_q.to(device="cuda") + + def _mask(cu_seqlens: torch.Tensor, max_seqlen: int) -> torch.Tensor: + # cu_seqlens may be longer than batch_size + 1 -- inference allocates it + # for the maximum batch size -- so only its first batch_size + 1 entries + # describe the current batch. + seqlens = cu_seqlens[1 : batch_size + 1] - cu_seqlens[:batch_size] + positions = torch.arange(max_seqlen, device=cu_seqlens.device) + # True marks a padding token, i.e. one beyond the sequence length. + return (positions.unsqueeze(0) >= seqlens.unsqueeze(1)).view(batch_size, 1, 1, max_seqlen) + + attention_mask_q = _mask(cu_seqlens_q, max_seqlen_q) if attention_type == "self": - attention_mask = attention_mask_q - else: - attention_mask = ( - attention_mask_q, - attention_mask_kv.to(device="cuda"), - ) - return attention_mask + return attention_mask_q + assert ( + cu_seqlens_kv is not None and max_seqlen_kv is not None + ), "cu_seqlens_kv and max_seqlen_kv are required for cross-attention" + return attention_mask_q, _mask(cu_seqlens_kv, max_seqlen_kv) @torch.no_grad() From e82b29d8d3ed6c4995ee513ea312874e820019a1 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 28 Jul 2026 17:05:56 +0200 Subject: [PATCH 2/2] Keep the padding mask on the attention device cu_seqlens is not required to be a CUDA tensor, and the previous implementation always moved the finished mask to CUDA. Building it on cu_seqlens' device regressed bshd/sbhd callers that pass CPU cu_seqlens: the mask reached the attention scores on the wrong device. Signed-off-by: Pawel Gadzinski --- .../pytorch/attention/dot_product_attention/utils.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 3e728cfa4a..8a4c6ce286 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1667,8 +1667,14 @@ def _mask(cu_seqlens: torch.Tensor, max_seqlen: int) -> torch.Tensor: # describe the current batch. seqlens = cu_seqlens[1 : batch_size + 1] - cu_seqlens[:batch_size] positions = torch.arange(max_seqlen, device=cu_seqlens.device) - # True marks a padding token, i.e. one beyond the sequence length. - return (positions.unsqueeze(0) >= seqlens.unsqueeze(1)).view(batch_size, 1, 1, max_seqlen) + # True marks a padding token, i.e. one beyond the sequence length. The + # mask is applied to the attention scores, so it goes on the device + # those live on -- a no-op unless cu_seqlens is a CPU tensor. + return ( + (positions.unsqueeze(0) >= seqlens.unsqueeze(1)) + .view(batch_size, 1, 1, max_seqlen) + .to(device="cuda") + ) attention_mask_q = _mask(cu_seqlens_q, max_seqlen_q) if attention_type == "self":