From 0d28d50f0c501da31120ad0a3745bfc129d90ffc Mon Sep 17 00:00:00 2001 From: Desh Raj Date: Wed, 29 Jul 2026 08:05:35 -0700 Subject: [PATCH 1/5] Add per-sequence causal policy to THD attention Signed-off-by: Desh Raj --- .../attention/test_mixed_thd_attention.py | 197 +++++++++++++++ .../dot_product_attention.py | 233 ++++++++++++++++++ 2 files changed, 430 insertions(+) create mode 100644 tests/pytorch/attention/test_mixed_thd_attention.py diff --git a/tests/pytorch/attention/test_mixed_thd_attention.py b/tests/pytorch/attention/test_mixed_thd_attention.py new file mode 100644 index 0000000000..dc6fb9b847 --- /dev/null +++ b/tests/pytorch/attention/test_mixed_thd_attention.py @@ -0,0 +1,197 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Tests for mixed causal/non-causal packed THD attention.""" + +import pytest +import torch + +from transformer_engine.pytorch import DotProductAttention + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + +NUM_HEADS = 16 +HEAD_DIM = 64 + + +def _make_dpa(dtype: torch.dtype) -> DotProductAttention: + return DotProductAttention( + num_attention_heads=NUM_HEADS, + kv_channels=HEAD_DIM, + attention_dropout=0.0, + qkv_format="thd", + attn_mask_type="padding", + tp_size=1, + tp_group=None, + layer_number=1, + ).to(device="cuda", dtype=dtype) + + +def _per_sequence_scalar_reference( + attention: DotProductAttention, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + cu_seqlens: torch.Tensor, + sequence_is_causal: torch.Tensor, +) -> torch.Tensor: + """Reference every packed sequence through its scalar-policy DPA call.""" + outputs = [] + start = 0 + for length, is_causal in zip( + (cu_seqlens[1:] - cu_seqlens[:-1]).tolist(), sequence_is_causal.tolist() + ): + end = start + length + sequence_cu_seqlens = torch.tensor((0, length), dtype=torch.int32, device=query.device) + outputs.append( + attention( + # THD scalar attention requires tensors with a zero storage + # offset. clone() is differentiable, so the reference still + # accumulates gradients into the original packed inputs. + query[start:end].clone(), + key[start:end].clone(), + value[start:end].clone(), + qkv_format="thd", + cu_seqlens_q=sequence_cu_seqlens, + cu_seqlens_kv=sequence_cu_seqlens, + max_seqlen_q=length, + max_seqlen_kv=length, + attn_mask_type="padding_causal" if is_causal else "padding", + window_size=(-1, 0) if is_causal else (-1, -1), + bottom_right_diagonal=False, + ) + ) + start = end + + return torch.cat(outputs, dim=0) + + +@pytest.mark.parametrize( + "sequence_is_causal", + ( + (False, True, False, True), + (False, False, False, False), + (True, True, True, True), + ), +) +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_mixed_thd_matches_scalar_forward_and_backward(sequence_is_causal, dtype): + """Mixed THD dispatch must match attention applied to each packed sequence.""" + if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported(): + pytest.skip("BF16 is not supported by this GPU") + + torch.manual_seed(1234) + lengths = torch.tensor((7, 3, 5, 4), dtype=torch.int32, device="cuda") + cu_seqlens = torch.cat((lengths.new_zeros(1), torch.cumsum(lengths, dim=0, dtype=torch.int32))) + token_count = int(lengths.sum().item()) + sequence_is_causal = torch.tensor(sequence_is_causal, dtype=torch.bool, device="cuda") + + def make_input(): + return ( + 0.1 * torch.randn(token_count, NUM_HEADS, HEAD_DIM, dtype=dtype, device="cuda") + ).requires_grad_() + + query, key, value = make_input(), make_input(), make_input() + reference_query = query.detach().clone().requires_grad_() + reference_key = key.detach().clone().requires_grad_() + reference_value = value.detach().clone().requires_grad_() + + attention = _make_dpa(dtype) + mixed_output = attention( + query, + key, + value, + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=int(lengths.max().item()), + max_seqlen_kv=int(lengths.max().item()), + thd_sequence_is_causal=sequence_is_causal, + ) + reference_output = _per_sequence_scalar_reference( + attention, + reference_query, + reference_key, + reference_value, + cu_seqlens, + sequence_is_causal, + ) + + tolerances = {"atol": 1.0e-3, "rtol": 1.0e-3} + if dtype == torch.bfloat16: + tolerances = {"atol": 1.5e-2, "rtol": 1.5e-2} + torch.testing.assert_close(mixed_output, reference_output, **tolerances) + + output_grad = torch.randn_like(mixed_output) + mixed_output.backward(output_grad) + reference_output.backward(output_grad) + for mixed_grad, reference_grad in ( + (query.grad, reference_query.grad), + (key.grad, reference_key.grad), + (value.grad, reference_value.grad), + ): + torch.testing.assert_close(mixed_grad, reference_grad, **tolerances) + + +@pytest.mark.parametrize( + ("invalid_case", "error"), + ( + ("policy_dtype", "one-dimensional torch.bool"), + ("sliding_window", "sliding-window attention"), + ), +) +def test_mixed_thd_rejects_unsupported_inputs(invalid_case, error): + """Reject metadata that the grouping path cannot preserve.""" + query = torch.randn(4, NUM_HEADS, HEAD_DIM, device="cuda", dtype=torch.float16) + key = torch.randn_like(query) + value = torch.randn_like(query) + cu_seqlens = torch.tensor((0, 2, 4), dtype=torch.int32, device="cuda") + attention = _make_dpa(torch.float16) + + kwargs = { + "thd_sequence_is_causal": torch.tensor((False, True), dtype=torch.bool, device="cuda"), + } + if invalid_case == "policy_dtype": + kwargs["thd_sequence_is_causal"] = torch.tensor((0, 1), dtype=torch.int32, device="cuda") + else: + kwargs["window_size"] = (4, 0) + + with pytest.raises(ValueError, match=error): + attention( + query, + key, + value, + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=2, + max_seqlen_kv=2, + **kwargs, + ) + + +def test_mixed_thd_empty_batch_preserves_input_gradients(): + """An all-empty packed batch needs zero gradients, not a disconnected output.""" + attention = _make_dpa(torch.float16) + query = torch.empty( + 0, NUM_HEADS, HEAD_DIM, device="cuda", dtype=torch.float16, requires_grad=True + ) + key = torch.empty_like(query, requires_grad=True) + value = torch.empty_like(query, requires_grad=True) + cu_seqlens = torch.zeros(1, dtype=torch.int32, device="cuda") + + output = attention( + query, + key, + value, + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=0, + max_seqlen_kv=0, + thd_sequence_is_causal=torch.empty(0, dtype=torch.bool, device="cuda"), + ) + assert output.shape == (0, NUM_HEADS * HEAD_DIM) + output.sum().backward() + for input_tensor in (query, key, value): + torch.testing.assert_close(input_tensor.grad, torch.zeros_like(input_tensor)) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index e6f50eb0da..a6137081c0 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1323,6 +1323,99 @@ def get_quantizer_roles( ] return base[:num_quantizers] + def _forward_mixed_thd( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + cu_seqlens: torch.Tensor, + thd_sequence_is_causal: torch.Tensor, + max_seqlen_q: Optional[int], + max_seqlen_kv: Optional[int], + *, + fast_zero_fill: bool, + ) -> torch.Tensor: + """Dispatch mixed causal/non-causal THD self-attention. + + Group whole sequences by policy, call the existing scalar-policy path + for each non-empty group, and restore the original packed token order. + """ + sequence_lengths = cu_seqlens[1:] - cu_seqlens[:-1] + token_count = query_layer.shape[0] + output_features = query_layer.shape[-2] * value_layer.shape[-1] + + if token_count == 0: + empty_output = query_layer.new_empty((0, output_features)) + return empty_output + 0 * (query_layer.sum() + key_layer.sum() + value_layer.sum()) + + def run_scalar_policy( + group_query: torch.Tensor, + group_key: torch.Tensor, + group_value: torch.Tensor, + group_cu_seqlens: torch.Tensor, + group_max_seqlen_q: Optional[int], + group_max_seqlen_kv: Optional[int], + is_causal: bool, + ) -> torch.Tensor: + return self.forward( + group_query, + group_key, + group_value, + qkv_format="thd", + cu_seqlens_q=group_cu_seqlens, + cu_seqlens_kv=group_cu_seqlens, + max_seqlen_q=group_max_seqlen_q, + max_seqlen_kv=group_max_seqlen_kv, + attn_mask_type="padding_causal" if is_causal else "padding", + window_size=(-1, 0) if is_causal else (-1, -1), + bottom_right_diagonal=False, + fast_zero_fill=fast_zero_fill, + ) + + first_policy = bool(thd_sequence_is_causal[0].item()) + if bool(torch.all(thd_sequence_is_causal == first_policy).item()): + return run_scalar_policy( + query_layer, + key_layer, + value_layer, + cu_seqlens, + max_seqlen_q, + max_seqlen_kv, + is_causal=first_policy, + ) + + token_is_causal = torch.repeat_interleave( + thd_sequence_is_causal, + sequence_lengths.to(dtype=torch.long), + output_size=token_count, + ) + output = query_layer.new_empty((token_count, output_features)) + for is_causal in (False, True): + sequence_mask = thd_sequence_is_causal == is_causal + group_token_indices = torch.nonzero(token_is_causal == is_causal).flatten() + if group_token_indices.numel() == 0: + continue + group_lengths = sequence_lengths[sequence_mask] + group_cu_seqlens = torch.cat( + ( + group_lengths.new_zeros(1), + torch.cumsum(group_lengths, dim=0, dtype=torch.int32), + ) + ) + group_max_seqlen = int(group_lengths.max().item()) + group_output = run_scalar_policy( + query_layer.index_select(0, group_token_indices), + key_layer.index_select(0, group_token_indices), + value_layer.index_select(0, group_token_indices), + group_cu_seqlens, + group_max_seqlen, + group_max_seqlen, + is_causal=is_causal, + ) + output = output.index_copy(0, group_token_indices, group_output) + + return output + @no_torch_dynamo(recursive=False) def forward( self, @@ -1357,6 +1450,7 @@ def forward( qkv_layer: Optional[torch.Tensor] = None, kv_layer: Optional[torch.Tensor] = None, qkv_interleave_dim: int = -3, + thd_sequence_is_causal: Optional[torch.Tensor] = None, ) -> torch.Tensor: r""" Dot Product Attention Layer. @@ -1596,8 +1690,147 @@ def forward( interleave sits; must be -3 (e.g. ``bs3hd``) or -2 (e.g. ``bsh3d``, Megatron-style). This is an explicit knob rather than shape inference, since e.g. ``h == 3`` would make the shapes ambiguous. + thd_sequence_is_causal: Optional[torch.Tensor], default = None + Per-sequence causal policy for packed THD self-attention. When supplied, + it must be a CUDA ``torch.bool`` tensor of shape ``[batch_size]``, where + ``batch_size = cu_seqlens_q.numel() - 1``. ``True`` selects causal + attention for that sequence and ``False`` selects full-context + attention. Transformer Engine groups the two policies and launches the + existing scalar-mask attention path once per non-empty group, then + restores the original packed token order. It therefore shares the + surrounding encoder computation but is not a single fused-attention + kernel. Initial support is limited to FP16/BF16 THD self-attention + without attention dropout, context parallelism, cacheing, FP8, + arbitrary padding between packed sequences, bias/ALiBi, score-mod, or + attention checkpointing. """ + if thd_sequence_is_causal is not None: + resolved_qkv_format = qkv_format if qkv_format is not None else self.qkv_format + if resolved_qkv_format != "thd": + raise ValueError("thd_sequence_is_causal is supported only with qkv_format='thd'.") + if ( + qkv_layer is not None + or kv_layer is not None + or query_layer is None + or key_layer is None + or value_layer is None + ): + raise ValueError( + "thd_sequence_is_causal requires separate query_layer, key_layer, and " + "value_layer." + ) + qkv = (query_layer, key_layer, value_layer) + if not all(tensor.is_cuda and tensor.device == query_layer.device for tensor in qkv): + raise ValueError( + "thd_sequence_is_causal requires CUDA Q, K, and V on the same device." + ) + if query_layer.dtype not in (torch.float16, torch.bfloat16) or not all( + tensor.dtype == query_layer.dtype for tensor in qkv + ): + raise ValueError( + "thd_sequence_is_causal requires FP16 or BF16 Q, K, and V of the same dtype." + ) + if self.attention_type != "self": + raise ValueError("thd_sequence_is_causal is supported only for self-attention.") + if cu_seqlens_q is None or cu_seqlens_kv is None: + raise ValueError("thd_sequence_is_causal requires cu_seqlens_q and cu_seqlens_kv.") + if not all( + cu_seqlens.ndim == 1 + and cu_seqlens.numel() >= 1 + and cu_seqlens.dtype == torch.int32 + and cu_seqlens.device == query_layer.device + for cu_seqlens in (cu_seqlens_q, cu_seqlens_kv) + ): + raise ValueError( + "thd_sequence_is_causal requires one-dimensional int32 cu_seqlens " + "on the Q/K/V device." + ) + if not torch.equal(cu_seqlens_q, cu_seqlens_kv): + raise ValueError( + "thd_sequence_is_causal requires identical Q and KV sequence boundaries." + ) + if thd_sequence_is_causal.dtype != torch.bool or thd_sequence_is_causal.ndim != 1: + raise ValueError( + "thd_sequence_is_causal must be a one-dimensional torch.bool tensor." + ) + if thd_sequence_is_causal.device != query_layer.device: + raise ValueError( + "thd_sequence_is_causal must be on the same device as query_layer." + ) + if thd_sequence_is_causal.numel() != cu_seqlens_q.numel() - 1: + raise ValueError( + "thd_sequence_is_causal must have one value per packed sequence: " + f"got {thd_sequence_is_causal.numel()} values for " + f"{cu_seqlens_q.numel() - 1} sequences." + ) + if not all( + tensor.ndim == 3 and tensor.shape[0] == query_layer.shape[0] for tensor in qkv + ): + raise ValueError( + "thd_sequence_is_causal requires THD Q, K, and V with the same token count." + ) + if int(cu_seqlens_q[-1].item()) != query_layer.shape[0]: + raise ValueError("The final cu_seqlens value must equal the THD Q/K/V token count.") + unsupported_options = ( + attention_mask is not None, + cu_seqlens_q_padded is not None, + cu_seqlens_kv_padded is not None, + inference_params is not None, + checkpoint_core_attention, + core_attention_bias_type != "no_bias", + core_attention_bias is not None, + alibi_slopes is not None, + score_mod is not None, + score_mod_bprop is not None, + score_mod_tensors is not None, + score_mod_bprop_tensors is not None, + fp8_output, + self.fp8, + self.attention_dropout != 0.0, + self.softmax_type != "vanilla", + self.return_max_logit, + self.cp_group is not None, + pad_between_seqs not in (None, False), + num_splits not in (None, 1), + ) + if any(unsupported_options): + raise ValueError( + "thd_sequence_is_causal currently supports only plain THD self-attention " + "without padding gaps, dropout, bias, score-mod, caching, FP8 output, " + "context parallelism, softmax variants, or attention checkpointing." + ) + if FP8GlobalStateManager.is_fp8_enabled() or FP8GlobalStateManager.is_fp8_calibration(): + raise ValueError("thd_sequence_is_causal does not yet support FP8 attention.") + if is_graph_capturing() or is_in_onnx_export_mode(): + raise ValueError( + "thd_sequence_is_causal does not yet support CUDA graph capture or ONNX export." + ) + effective_window_size = self.window_size if window_size is None else window_size + if effective_window_size not in {(-1, -1), (-1, 0)}: + raise ValueError( + "thd_sequence_is_causal does not yet support sliding-window attention." + ) + effective_bottom_right_diagonal = ( + self.bottom_right_diagonal + if bottom_right_diagonal is None + else bottom_right_diagonal + ) + if effective_bottom_right_diagonal is True: + raise ValueError( + "thd_sequence_is_causal requires the standard top-left causal diagonal." + ) + return self._forward_mixed_thd( + query_layer, + key_layer, + value_layer, + cu_seqlens_q, + thd_sequence_is_causal, + max_seqlen_q, + max_seqlen_kv, + fast_zero_fill=fast_zero_fill, + ) + query_layer, key_layer, value_layer, declared_qkv_layout = _unpack_packed_qkv( qkv_layer, kv_layer, From 50d2c81201a33a31521faab954a98e412367b614 Mon Sep 17 00:00:00 2001 From: Desh Raj Date: Thu, 30 Jul 2026 09:29:20 -0700 Subject: [PATCH 2/5] Fix mixed THD test CI coverage Signed-off-by: Desh Raj --- qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/attention/test_mixed_thd_attention.py | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 759432857a..f0de413133 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -63,6 +63,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_att python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_uproj.xml $TE_PATH/tests/pytorch/attention/test_fused_mla_q_uproj.py || test_fail "test_fused_mla_q_uproj.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mixed_thd_attention.xml $TE_PATH/tests/pytorch/attention/test_mixed_thd_attention.py || test_fail "test_mixed_thd_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" export NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint if [ ! -d "$NVTE_TEST_CHECKPOINT_ARTIFACT_PATH" ]; then diff --git a/tests/pytorch/attention/test_mixed_thd_attention.py b/tests/pytorch/attention/test_mixed_thd_attention.py index dc6fb9b847..04624eece4 100644 --- a/tests/pytorch/attention/test_mixed_thd_attention.py +++ b/tests/pytorch/attention/test_mixed_thd_attention.py @@ -1,4 +1,6 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. """Tests for mixed causal/non-causal packed THD attention.""" From a88c9b64b291bccea2373f15a9bcf90626508afb Mon Sep 17 00:00:00 2001 From: Desh Raj Date: Fri, 7 Aug 2026 10:53:39 -0700 Subject: [PATCH 3/5] Use inter-sequence padding for mixed THD masks Replace the per-policy Q/K/V gather and output scatter dispatcher with the review-suggested inter-sequence-padding design. Keep the original packed Q/K/V storage, derive logical and physical cu-seqlens for each mask policy, run the ordinary scalar-mask DPA path with pad_between_seqs, and sum its disjoint zero-padded outputs. Expose a generic attn_mask_type_per_seq mapping plus window_size_per_mask_type so callers can combine full-context offline attention with bounded causal attention. Preserve the scalar fast path for uniform batches, avoid policy-dependent host synchronization, and cover forward/backward parity, existing physical padding, cross-attention, per-policy windows, bottom-right masks, invalid inputs, empty batches, and CUDA graph replay. NRT H100 validation passed all 16 TransformerEngine cases and 3 focused Megatron integration tests. In an 8-node/64-H100 matched-auto comparison over iterations 7-100, grouped averaged 1802.289 ms and pad-between averaged 1779.596 ms (-1.259%); the paired 95% interval crossed zero, indicating practical parity when the backend is held constant. In the deployable-backend comparison, grouped/Flash averaged 1596.129 ms while pad-between/auto-cuDNN averaged 1766.769 ms: pad-between was 10.691% slower by mean and 11.482% slower by 5%-trimmed mean. The paired penalty was +170.640 ms with a 95% interval of +89.728 to +251.552 ms, with all 16 fully logged workload counters matching across 100 steps. The implementation removes the original data-movement contention, but the current image cannot run full-context inter-sequence padding through FlashAttention 2 and has no FA3, so the path falls back to cuDNN fused attention. Retain grouped/Flash for production until a competitive padding-capable backend is available; this commit keeps the cleaner generic API for review and future backend support. Signed-off-by: Desh Raj --- .../attention/test_mixed_thd_attention.py | 318 ++++++++++--- .../dot_product_attention.py | 446 ++++++++++-------- 2 files changed, 495 insertions(+), 269 deletions(-) diff --git a/tests/pytorch/attention/test_mixed_thd_attention.py b/tests/pytorch/attention/test_mixed_thd_attention.py index 04624eece4..4cb3f39578 100644 --- a/tests/pytorch/attention/test_mixed_thd_attention.py +++ b/tests/pytorch/attention/test_mixed_thd_attention.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""Tests for mixed causal/non-causal packed THD attention.""" +"""Tests for per-sequence mask types in packed THD attention.""" import pytest import torch @@ -16,77 +16,110 @@ HEAD_DIM = 64 -def _make_dpa(dtype: torch.dtype) -> DotProductAttention: +def _make_dpa(dtype: torch.dtype, attention_type: str = "self") -> DotProductAttention: return DotProductAttention( num_attention_heads=NUM_HEADS, kv_channels=HEAD_DIM, attention_dropout=0.0, qkv_format="thd", attn_mask_type="padding", + attention_type=attention_type, tp_size=1, tp_group=None, layer_number=1, ).to(device="cuda", dtype=dtype) +def _make_cu_seqlens(lengths): + lengths = torch.tensor(lengths, dtype=torch.int32, device="cuda") + return torch.cat((lengths.new_zeros(1), torch.cumsum(lengths, dim=0, dtype=torch.int32))) + + +def _make_policy(sequence_ids): + return { + mask_type: torch.tensor(ids, dtype=torch.int64, device="cuda") + for mask_type, ids in sequence_ids.items() + } + + def _per_sequence_scalar_reference( attention: DotProductAttention, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - cu_seqlens: torch.Tensor, - sequence_is_causal: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + mask_type_per_seq, + *, + cu_seqlens_q_padded=None, + cu_seqlens_kv_padded=None, + window_size_per_mask_type=None, ) -> torch.Tensor: - """Reference every packed sequence through its scalar-policy DPA call.""" - outputs = [] - start = 0 - for length, is_causal in zip( - (cu_seqlens[1:] - cu_seqlens[:-1]).tolist(), sequence_is_causal.tolist() - ): - end = start + length - sequence_cu_seqlens = torch.tensor((0, length), dtype=torch.int32, device=query.device) - outputs.append( - attention( - # THD scalar attention requires tensors with a zero storage - # offset. clone() is differentiable, so the reference still - # accumulates gradients into the original packed inputs. - query[start:end].clone(), - key[start:end].clone(), - value[start:end].clone(), + """Reference each physical sequence through an ordinary scalar-mask call.""" + q_physical = cu_seqlens_q if cu_seqlens_q_padded is None else cu_seqlens_q_padded + kv_physical = cu_seqlens_kv if cu_seqlens_kv_padded is None else cu_seqlens_kv_padded + outputs = [None] * (cu_seqlens_q.numel() - 1) + for mask_type, sequence_ids in mask_type_per_seq.items(): + for sequence_id in sequence_ids.tolist(): + q_length = int((cu_seqlens_q[sequence_id + 1] - cu_seqlens_q[sequence_id]).item()) + kv_length = int( + (cu_seqlens_kv[sequence_id + 1] - cu_seqlens_kv[sequence_id]).item() + ) + q_start = int(q_physical[sequence_id].item()) + kv_start = int(kv_physical[sequence_id].item()) + q_end = q_start + q_length + kv_end = kv_start + kv_length + sequence_cu_seqlens_q = torch.tensor( + (0, q_length), dtype=torch.int32, device=query.device + ) + sequence_cu_seqlens_kv = torch.tensor( + (0, kv_length), dtype=torch.int32, device=query.device + ) + policy_window_size = None + if window_size_per_mask_type is not None: + policy_window_size = window_size_per_mask_type.get(mask_type) + outputs[sequence_id] = attention( + # THD scalar attention requires zero-offset storage. clone() is + # differentiable, so gradients still reach the packed inputs. + query[q_start:q_end].clone(), + key[kv_start:kv_end].clone(), + value[kv_start:kv_end].clone(), qkv_format="thd", - cu_seqlens_q=sequence_cu_seqlens, - cu_seqlens_kv=sequence_cu_seqlens, - max_seqlen_q=length, - max_seqlen_kv=length, - attn_mask_type="padding_causal" if is_causal else "padding", - window_size=(-1, 0) if is_causal else (-1, -1), - bottom_right_diagonal=False, + cu_seqlens_q=sequence_cu_seqlens_q, + cu_seqlens_kv=sequence_cu_seqlens_kv, + max_seqlen_q=q_length, + max_seqlen_kv=kv_length, + attn_mask_type=mask_type, + window_size=policy_window_size, ) - ) - start = end - return torch.cat(outputs, dim=0) + output = query.new_zeros((query.shape[0], NUM_HEADS * HEAD_DIM)) + for sequence_id, sequence_output in enumerate(outputs): + q_length = int((cu_seqlens_q[sequence_id + 1] - cu_seqlens_q[sequence_id]).item()) + q_start = int(q_physical[sequence_id].item()) + output[q_start : q_start + q_length] = sequence_output + return output @pytest.mark.parametrize( - "sequence_is_causal", + "mask_type_per_seq", ( - (False, True, False, True), - (False, False, False, False), - (True, True, True, True), + {"padding": (0, 2), "padding_causal": (1, 3)}, + {"padding": (0, 1, 2, 3)}, + {"padding_causal": (0, 1, 2, 3)}, + {"padding": (0, 2), "padding_causal_bottom_right": (1, 3)}, ), ) @pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) -def test_mixed_thd_matches_scalar_forward_and_backward(sequence_is_causal, dtype): - """Mixed THD dispatch must match attention applied to each packed sequence.""" +def test_thd_mask_types_match_scalar_forward_and_backward(mask_type_per_seq, dtype): + """Per-sequence masks must match independent scalar-policy attention.""" if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported(): pytest.skip("BF16 is not supported by this GPU") torch.manual_seed(1234) - lengths = torch.tensor((7, 3, 5, 4), dtype=torch.int32, device="cuda") - cu_seqlens = torch.cat((lengths.new_zeros(1), torch.cumsum(lengths, dim=0, dtype=torch.int32))) - token_count = int(lengths.sum().item()) - sequence_is_causal = torch.tensor(sequence_is_causal, dtype=torch.bool, device="cuda") + cu_seqlens = _make_cu_seqlens((7, 3, 5, 4)) + token_count = 19 + mask_type_per_seq = _make_policy(mask_type_per_seq) def make_input(): return ( @@ -99,16 +132,16 @@ def make_input(): reference_value = value.detach().clone().requires_grad_() attention = _make_dpa(dtype) - mixed_output = attention( + output = attention( query, key, value, qkv_format="thd", cu_seqlens_q=cu_seqlens, cu_seqlens_kv=cu_seqlens, - max_seqlen_q=int(lengths.max().item()), - max_seqlen_kv=int(lengths.max().item()), - thd_sequence_is_causal=sequence_is_causal, + max_seqlen_q=7, + max_seqlen_kv=7, + attn_mask_type_per_seq=mask_type_per_seq, ) reference_output = _per_sequence_scalar_reference( attention, @@ -116,47 +149,205 @@ def make_input(): reference_key, reference_value, cu_seqlens, - sequence_is_causal, + cu_seqlens, + mask_type_per_seq, ) tolerances = {"atol": 1.0e-3, "rtol": 1.0e-3} if dtype == torch.bfloat16: tolerances = {"atol": 1.5e-2, "rtol": 1.5e-2} - torch.testing.assert_close(mixed_output, reference_output, **tolerances) + torch.testing.assert_close(output, reference_output, **tolerances) - output_grad = torch.randn_like(mixed_output) - mixed_output.backward(output_grad) + output_grad = torch.randn_like(output) + output.backward(output_grad) reference_output.backward(output_grad) - for mixed_grad, reference_grad in ( + for actual_grad, reference_grad in ( (query.grad, reference_query.grad), (key.grad, reference_key.grad), (value.grad, reference_value.grad), ): - torch.testing.assert_close(mixed_grad, reference_grad, **tolerances) + torch.testing.assert_close(actual_grad, reference_grad, **tolerances) + + +def test_thd_mask_types_compose_with_existing_inter_sequence_padding(): + """Policy gaps and pre-existing THD padding must share the physical layout.""" + torch.manual_seed(1234) + dtype = torch.float16 + cu_seqlens = _make_cu_seqlens((7, 3, 5, 4)) + cu_seqlens_padded = torch.tensor((0, 9, 12, 19, 23), dtype=torch.int32, device="cuda") + mask_type_per_seq = _make_policy({"padding": (0, 2), "padding_causal": (1, 3)}) + + def make_input(): + return ( + 0.1 * torch.randn(23, NUM_HEADS, HEAD_DIM, dtype=dtype, device="cuda") + ).requires_grad_() + + query, key, value = make_input(), make_input(), make_input() + reference_query = query.detach().clone().requires_grad_() + reference_key = key.detach().clone().requires_grad_() + reference_value = value.detach().clone().requires_grad_() + attention = _make_dpa(dtype) + + output = attention( + query, + key, + value, + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=7, + max_seqlen_kv=7, + attn_mask_type_per_seq=mask_type_per_seq, + ) + reference_output = _per_sequence_scalar_reference( + attention, + reference_query, + reference_key, + reference_value, + cu_seqlens, + cu_seqlens, + mask_type_per_seq, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + ) + torch.testing.assert_close(output, reference_output, atol=1.0e-3, rtol=1.0e-3) + + output_grad = torch.randn_like(output) + output.backward(output_grad) + reference_output.backward(output_grad) + for actual_grad, reference_grad in ( + (query.grad, reference_query.grad), + (key.grad, reference_key.grad), + (value.grad, reference_value.grad), + ): + torch.testing.assert_close(actual_grad, reference_grad, atol=1.0e-3, rtol=1.0e-3) + + +def test_thd_mask_types_support_cross_attention_and_per_policy_windows(): + """Mask groups may use different Q/KV lengths and sliding windows.""" + torch.manual_seed(1234) + dtype = torch.float16 + cu_seqlens_q = _make_cu_seqlens((5, 3, 4)) + cu_seqlens_kv = _make_cu_seqlens((7, 2, 5)) + mask_type_per_seq = _make_policy( + {"padding": (0, 2), "padding_causal_bottom_right": (1,)} + ) + windows = {"padding": (-1, -1), "padding_causal_bottom_right": (2, 0)} + + query = ( + 0.1 * torch.randn(12, NUM_HEADS, HEAD_DIM, dtype=dtype, device="cuda") + ).requires_grad_() + key = (0.1 * torch.randn(14, NUM_HEADS, HEAD_DIM, dtype=dtype, device="cuda")).requires_grad_() + value = (0.1 * torch.randn_like(key)).requires_grad_() + reference_query = query.detach().clone().requires_grad_() + reference_key = key.detach().clone().requires_grad_() + reference_value = value.detach().clone().requires_grad_() + attention = _make_dpa(dtype, attention_type="cross") + + output = attention( + query, + key, + value, + qkv_format="thd", + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=5, + max_seqlen_kv=7, + attn_mask_type_per_seq=mask_type_per_seq, + window_size_per_mask_type=windows, + ) + reference_output = _per_sequence_scalar_reference( + attention, + reference_query, + reference_key, + reference_value, + cu_seqlens_q, + cu_seqlens_kv, + mask_type_per_seq, + window_size_per_mask_type=windows, + ) + torch.testing.assert_close(output, reference_output, atol=1.0e-3, rtol=1.0e-3) + + output_grad = torch.randn_like(output) + output.backward(output_grad) + reference_output.backward(output_grad) + for actual_grad, reference_grad in ( + (query.grad, reference_query.grad), + (key.grad, reference_key.grad), + (value.grad, reference_value.grad), + ): + torch.testing.assert_close(actual_grad, reference_grad, atol=1.0e-3, rtol=1.0e-3) + + +def test_thd_mask_types_support_cuda_graph_capture(): + """The sync-free policy metadata path must be replayable in a CUDA graph.""" + torch.manual_seed(1234) + dtype = torch.float16 + cu_seqlens = _make_cu_seqlens((7, 3, 5, 4)) + mask_type_per_seq = _make_policy({"padding": (0, 2), "padding_causal": (1, 3)}) + query = 0.1 * torch.randn(19, NUM_HEADS, HEAD_DIM, dtype=dtype, device="cuda") + key = 0.1 * torch.randn_like(query) + value = 0.1 * torch.randn_like(query) + attention = _make_dpa(dtype).eval() + + def run_attention(): + return attention( + query, + key, + value, + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=7, + max_seqlen_kv=7, + attn_mask_type_per_seq=mask_type_per_seq, + ) + + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + for _ in range(3): + run_attention() + torch.cuda.current_stream().wait_stream(warmup_stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_output = run_attention() + graph.replay() + actual_output = captured_output.clone() + expected_output = run_attention() + torch.testing.assert_close(actual_output, expected_output, atol=1.0e-3, rtol=1.0e-3) @pytest.mark.parametrize( ("invalid_case", "error"), ( - ("policy_dtype", "one-dimensional torch.bool"), - ("sliding_window", "sliding-window attention"), + ("policy_dtype", "one-dimensional int32 or int64"), + ("missing_sequence", "assign every sequence exactly once by count"), + ("mask_type", "only padding mask types"), + ("window_key", "mask not present"), ), ) -def test_mixed_thd_rejects_unsupported_inputs(invalid_case, error): - """Reject metadata that the grouping path cannot preserve.""" +def test_thd_mask_types_reject_invalid_inputs(invalid_case, error): + """Reject malformed policy metadata before launching attention.""" query = torch.randn(4, NUM_HEADS, HEAD_DIM, device="cuda", dtype=torch.float16) key = torch.randn_like(query) value = torch.randn_like(query) cu_seqlens = torch.tensor((0, 2, 4), dtype=torch.int32, device="cuda") attention = _make_dpa(torch.float16) + policies = _make_policy({"padding": (0,), "padding_causal": (1,)}) + windows = None - kwargs = { - "thd_sequence_is_causal": torch.tensor((False, True), dtype=torch.bool, device="cuda"), - } if invalid_case == "policy_dtype": - kwargs["thd_sequence_is_causal"] = torch.tensor((0, 1), dtype=torch.int32, device="cuda") + policies["padding"] = torch.tensor((0,), dtype=torch.float32, device="cuda") + elif invalid_case == "missing_sequence": + policies = _make_policy({"padding": (0,)}) + elif invalid_case == "mask_type": + policies = _make_policy({"no_mask": (0,), "padding": (1,)}) else: - kwargs["window_size"] = (4, 0) + windows = {"padding_causal_bottom_right": (-1, 0)} with pytest.raises(ValueError, match=error): attention( @@ -168,12 +359,13 @@ def test_mixed_thd_rejects_unsupported_inputs(invalid_case, error): cu_seqlens_kv=cu_seqlens, max_seqlen_q=2, max_seqlen_kv=2, - **kwargs, + attn_mask_type_per_seq=policies, + window_size_per_mask_type=windows, ) -def test_mixed_thd_empty_batch_preserves_input_gradients(): - """An all-empty packed batch needs zero gradients, not a disconnected output.""" +def test_thd_mask_types_empty_batch_preserves_input_gradients(): + """An empty packed batch needs zero gradients, not a disconnected output.""" attention = _make_dpa(torch.float16) query = torch.empty( 0, NUM_HEADS, HEAD_DIM, device="cuda", dtype=torch.float16, requires_grad=True @@ -191,7 +383,7 @@ def test_mixed_thd_empty_batch_preserves_input_gradients(): cu_seqlens_kv=cu_seqlens, max_seqlen_q=0, max_seqlen_kv=0, - thd_sequence_is_causal=torch.empty(0, dtype=torch.bool, device="cuda"), + attn_mask_type_per_seq={"padding": torch.empty(0, dtype=torch.int64, device="cuda")}, ) assert output.shape == (0, NUM_HEADS * HEAD_DIM) output.sum().backward() diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index a6137081c0..4d286852cb 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1323,97 +1323,135 @@ def get_quantizer_roles( ] return base[:num_quantizers] - def _forward_mixed_thd( + @staticmethod + def _select_thd_sequences( + cu_seqlens: torch.Tensor, + cu_seqlens_padded: Optional[torch.Tensor], + sequence_ids: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Describe selected THD sequences without moving their token storage.""" + sequence_ids = sequence_ids.to(dtype=torch.long) + physical_cu_seqlens = cu_seqlens if cu_seqlens_padded is None else cu_seqlens_padded + physical_starts = physical_cu_seqlens.index_select(0, sequence_ids) + sequence_lengths = cu_seqlens.index_select(0, sequence_ids + 1) - cu_seqlens.index_select( + 0, sequence_ids + ) + selected_cu_seqlens = torch.cat( + ( + sequence_lengths.new_zeros(1), + torch.cumsum(sequence_lengths, dim=0, dtype=torch.int32), + ) + ) + selected_cu_seqlens_padded = torch.cat( + (physical_starts, physical_cu_seqlens[-1:]) + ) + return selected_cu_seqlens, selected_cu_seqlens_padded + + def _forward_thd_mask_types( self, query_layer: torch.Tensor, key_layer: torch.Tensor, value_layer: torch.Tensor, - cu_seqlens: torch.Tensor, - thd_sequence_is_causal: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + attn_mask_type_per_seq: Dict[str, torch.Tensor], + window_size_per_mask_type: Optional[Dict[str, Tuple[int, int]]], max_seqlen_q: Optional[int], max_seqlen_kv: Optional[int], *, + window_size: Optional[Tuple[int, int]], + checkpoint_core_attention: bool, fast_zero_fill: bool, + pad_between_seqs: Optional[bool], + bf16_backward: Optional[bool], + num_splits: Optional[int], ) -> torch.Tensor: - """Dispatch mixed causal/non-causal THD self-attention. + """Dispatch per-sequence THD masks without regrouping Q/K/V tokens. - Group whole sequences by policy, call the existing scalar-policy path - for each non-empty group, and restore the original packed token order. + Each scalar-mask call sees the original Q/K/V storage. Sequences assigned + to other mask types are represented as inter-sequence padding, so the + backend leaves zeros at those positions. The disjoint outputs can then be + added without a token gather or scatter. """ - sequence_lengths = cu_seqlens[1:] - cu_seqlens[:-1] - token_count = query_layer.shape[0] output_features = query_layer.shape[-2] * value_layer.shape[-1] - - if token_count == 0: + if query_layer.shape[0] == 0: empty_output = query_layer.new_empty((0, output_features)) return empty_output + 0 * (query_layer.sum() + key_layer.sum() + value_layer.sum()) - def run_scalar_policy( - group_query: torch.Tensor, - group_key: torch.Tensor, - group_value: torch.Tensor, - group_cu_seqlens: torch.Tensor, - group_max_seqlen_q: Optional[int], - group_max_seqlen_kv: Optional[int], - is_causal: bool, - ) -> torch.Tensor: + nonempty_mask_types = [ + mask_type + for mask_type, sequence_ids in attn_mask_type_per_seq.items() + if sequence_ids.numel() > 0 + ] + batch_size = cu_seqlens_q.shape[0] - 1 + if ( + len(nonempty_mask_types) == 1 + and attn_mask_type_per_seq[nonempty_mask_types[0]].numel() == batch_size + ): + mask_type = nonempty_mask_types[0] + policy_window_size = window_size + if window_size_per_mask_type is not None: + policy_window_size = window_size_per_mask_type.get(mask_type, window_size) return self.forward( - group_query, - group_key, - group_value, + query_layer, + key_layer, + value_layer, qkv_format="thd", - cu_seqlens_q=group_cu_seqlens, - cu_seqlens_kv=group_cu_seqlens, - max_seqlen_q=group_max_seqlen_q, - max_seqlen_kv=group_max_seqlen_kv, - attn_mask_type="padding_causal" if is_causal else "padding", - window_size=(-1, 0) if is_causal else (-1, -1), - bottom_right_diagonal=False, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + attn_mask_type=mask_type, + window_size=policy_window_size, + checkpoint_core_attention=checkpoint_core_attention, fast_zero_fill=fast_zero_fill, + pad_between_seqs=pad_between_seqs, + bf16_backward=bf16_backward, + num_splits=num_splits, ) - first_policy = bool(thd_sequence_is_causal[0].item()) - if bool(torch.all(thd_sequence_is_causal == first_policy).item()): - return run_scalar_policy( + output = None + for mask_type in nonempty_mask_types: + sequence_ids = attn_mask_type_per_seq[mask_type] + selected_cu_seqlens_q, selected_cu_seqlens_q_padded = self._select_thd_sequences( + cu_seqlens_q, + cu_seqlens_q_padded, + sequence_ids, + ) + selected_cu_seqlens_kv, selected_cu_seqlens_kv_padded = self._select_thd_sequences( + cu_seqlens_kv, + cu_seqlens_kv_padded, + sequence_ids, + ) + policy_window_size = window_size + if window_size_per_mask_type is not None: + policy_window_size = window_size_per_mask_type.get(mask_type, window_size) + policy_output = self.forward( query_layer, key_layer, value_layer, - cu_seqlens, - max_seqlen_q, - max_seqlen_kv, - is_causal=first_policy, - ) - - token_is_causal = torch.repeat_interleave( - thd_sequence_is_causal, - sequence_lengths.to(dtype=torch.long), - output_size=token_count, - ) - output = query_layer.new_empty((token_count, output_features)) - for is_causal in (False, True): - sequence_mask = thd_sequence_is_causal == is_causal - group_token_indices = torch.nonzero(token_is_causal == is_causal).flatten() - if group_token_indices.numel() == 0: - continue - group_lengths = sequence_lengths[sequence_mask] - group_cu_seqlens = torch.cat( - ( - group_lengths.new_zeros(1), - torch.cumsum(group_lengths, dim=0, dtype=torch.int32), - ) - ) - group_max_seqlen = int(group_lengths.max().item()) - group_output = run_scalar_policy( - query_layer.index_select(0, group_token_indices), - key_layer.index_select(0, group_token_indices), - value_layer.index_select(0, group_token_indices), - group_cu_seqlens, - group_max_seqlen, - group_max_seqlen, - is_causal=is_causal, + qkv_format="thd", + cu_seqlens_q=selected_cu_seqlens_q, + cu_seqlens_kv=selected_cu_seqlens_kv, + cu_seqlens_q_padded=selected_cu_seqlens_q_padded, + cu_seqlens_kv_padded=selected_cu_seqlens_kv_padded, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + attn_mask_type=mask_type, + window_size=policy_window_size, + checkpoint_core_attention=checkpoint_core_attention, + fast_zero_fill=fast_zero_fill, + pad_between_seqs=True, + bf16_backward=bf16_backward, + num_splits=num_splits, ) - output = output.index_copy(0, group_token_indices, group_output) - + output = policy_output if output is None else output + policy_output + if output is None: + raise ValueError("attn_mask_type_per_seq must assign at least one sequence.") return output @no_torch_dynamo(recursive=False) @@ -1450,7 +1488,8 @@ def forward( qkv_layer: Optional[torch.Tensor] = None, kv_layer: Optional[torch.Tensor] = None, qkv_interleave_dim: int = -3, - thd_sequence_is_causal: Optional[torch.Tensor] = None, + attn_mask_type_per_seq: Optional[Dict[str, torch.Tensor]] = None, + window_size_per_mask_type: Optional[Dict[str, Tuple[int, int]]] = None, ) -> torch.Tensor: r""" Dot Product Attention Layer. @@ -1690,146 +1729,59 @@ def forward( interleave sits; must be -3 (e.g. ``bs3hd``) or -2 (e.g. ``bsh3d``, Megatron-style). This is an explicit knob rather than shape inference, since e.g. ``h == 3`` would make the shapes ambiguous. - thd_sequence_is_causal: Optional[torch.Tensor], default = None - Per-sequence causal policy for packed THD self-attention. When supplied, - it must be a CUDA ``torch.bool`` tensor of shape ``[batch_size]``, where - ``batch_size = cu_seqlens_q.numel() - 1``. ``True`` selects causal - attention for that sequence and ``False`` selects full-context - attention. Transformer Engine groups the two policies and launches the - existing scalar-mask attention path once per non-empty group, then - restores the original packed token order. It therefore shares the - surrounding encoder computation but is not a single fused-attention - kernel. Initial support is limited to FP16/BF16 THD self-attention - without attention dropout, context parallelism, cacheing, FP8, - arbitrary padding between packed sequences, bias/ALiBi, score-mod, or - attention checkpointing. + attn_mask_type_per_seq: Optional[Dict[str, torch.Tensor]], default = None + Per-sequence mask policies for packed THD attention. Each key is an + attention mask type and each value is a one-dimensional, ascending CUDA + integer tensor containing the sequence IDs assigned to that mask. The + sequence-ID tensors must be disjoint and together contain every sequence + exactly once. Transformer Engine keeps Q/K/V in their original packed + layout and describes the sequences assigned to other policies as + inter-sequence padding. It launches the existing scalar-mask attention + path once per non-empty policy and adds the disjoint outputs. + window_size_per_mask_type: Optional[Dict[str, Tuple[int, int]]], default = None + Optional sliding-window override for each key in + :attr:`attn_mask_type_per_seq`. This is useful when, for example, offline + sequences need full context while causal sequences use bounded left + context. Missing keys fall back to :attr:`window_size`. """ - if thd_sequence_is_causal is not None: - resolved_qkv_format = qkv_format if qkv_format is not None else self.qkv_format - if resolved_qkv_format != "thd": - raise ValueError("thd_sequence_is_causal is supported only with qkv_format='thd'.") - if ( - qkv_layer is not None - or kv_layer is not None - or query_layer is None - or key_layer is None - or value_layer is None - ): - raise ValueError( - "thd_sequence_is_causal requires separate query_layer, key_layer, and " - "value_layer." - ) - qkv = (query_layer, key_layer, value_layer) - if not all(tensor.is_cuda and tensor.device == query_layer.device for tensor in qkv): - raise ValueError( - "thd_sequence_is_causal requires CUDA Q, K, and V on the same device." - ) - if query_layer.dtype not in (torch.float16, torch.bfloat16) or not all( - tensor.dtype == query_layer.dtype for tensor in qkv - ): - raise ValueError( - "thd_sequence_is_causal requires FP16 or BF16 Q, K, and V of the same dtype." - ) - if self.attention_type != "self": - raise ValueError("thd_sequence_is_causal is supported only for self-attention.") - if cu_seqlens_q is None or cu_seqlens_kv is None: - raise ValueError("thd_sequence_is_causal requires cu_seqlens_q and cu_seqlens_kv.") - if not all( - cu_seqlens.ndim == 1 - and cu_seqlens.numel() >= 1 - and cu_seqlens.dtype == torch.int32 - and cu_seqlens.device == query_layer.device - for cu_seqlens in (cu_seqlens_q, cu_seqlens_kv) - ): - raise ValueError( - "thd_sequence_is_causal requires one-dimensional int32 cu_seqlens " - "on the Q/K/V device." - ) - if not torch.equal(cu_seqlens_q, cu_seqlens_kv): - raise ValueError( - "thd_sequence_is_causal requires identical Q and KV sequence boundaries." - ) - if thd_sequence_is_causal.dtype != torch.bool or thd_sequence_is_causal.ndim != 1: - raise ValueError( - "thd_sequence_is_causal must be a one-dimensional torch.bool tensor." - ) - if thd_sequence_is_causal.device != query_layer.device: - raise ValueError( - "thd_sequence_is_causal must be on the same device as query_layer." - ) - if thd_sequence_is_causal.numel() != cu_seqlens_q.numel() - 1: - raise ValueError( - "thd_sequence_is_causal must have one value per packed sequence: " - f"got {thd_sequence_is_causal.numel()} values for " - f"{cu_seqlens_q.numel() - 1} sequences." - ) - if not all( - tensor.ndim == 3 and tensor.shape[0] == query_layer.shape[0] for tensor in qkv - ): - raise ValueError( - "thd_sequence_is_causal requires THD Q, K, and V with the same token count." - ) - if int(cu_seqlens_q[-1].item()) != query_layer.shape[0]: - raise ValueError("The final cu_seqlens value must equal the THD Q/K/V token count.") - unsupported_options = ( - attention_mask is not None, - cu_seqlens_q_padded is not None, - cu_seqlens_kv_padded is not None, - inference_params is not None, - checkpoint_core_attention, - core_attention_bias_type != "no_bias", - core_attention_bias is not None, - alibi_slopes is not None, - score_mod is not None, - score_mod_bprop is not None, - score_mod_tensors is not None, - score_mod_bprop_tensors is not None, - fp8_output, - self.fp8, - self.attention_dropout != 0.0, - self.softmax_type != "vanilla", - self.return_max_logit, - self.cp_group is not None, - pad_between_seqs not in (None, False), - num_splits not in (None, 1), - ) - if any(unsupported_options): - raise ValueError( - "thd_sequence_is_causal currently supports only plain THD self-attention " - "without padding gaps, dropout, bias, score-mod, caching, FP8 output, " - "context parallelism, softmax variants, or attention checkpointing." - ) - if FP8GlobalStateManager.is_fp8_enabled() or FP8GlobalStateManager.is_fp8_calibration(): - raise ValueError("thd_sequence_is_causal does not yet support FP8 attention.") - if is_graph_capturing() or is_in_onnx_export_mode(): - raise ValueError( - "thd_sequence_is_causal does not yet support CUDA graph capture or ONNX export." - ) - effective_window_size = self.window_size if window_size is None else window_size - if effective_window_size not in {(-1, -1), (-1, 0)}: - raise ValueError( - "thd_sequence_is_causal does not yet support sliding-window attention." - ) - effective_bottom_right_diagonal = ( - self.bottom_right_diagonal - if bottom_right_diagonal is None - else bottom_right_diagonal - ) - if effective_bottom_right_diagonal is True: - raise ValueError( - "thd_sequence_is_causal requires the standard top-left causal diagonal." - ) - return self._forward_mixed_thd( - query_layer, - key_layer, - value_layer, - cu_seqlens_q, - thd_sequence_is_causal, - max_seqlen_q, - max_seqlen_kv, - fast_zero_fill=fast_zero_fill, + if window_size_per_mask_type is not None and attn_mask_type_per_seq is None: + raise ValueError( + "window_size_per_mask_type requires attn_mask_type_per_seq." ) + if attn_mask_type_per_seq is not None: + if not isinstance(attn_mask_type_per_seq, dict) or not attn_mask_type_per_seq: + raise ValueError("attn_mask_type_per_seq must be a non-empty dictionary.") + normalized_mask_types = {} + for mask_type, sequence_ids in attn_mask_type_per_seq.items(): + if not isinstance(mask_type, str): + raise ValueError("attn_mask_type_per_seq keys must be attention mask strings.") + mask_type = mask_type.replace(",", "_") + if mask_type == "causal_padding": + mask_type = "padding_causal" + if mask_type not in AttnMaskTypes or "padding" not in mask_type: + raise ValueError( + "attn_mask_type_per_seq supports only padding mask types, got " + f"{mask_type!r}." + ) + if mask_type in normalized_mask_types: + raise ValueError(f"Duplicate normalized attention mask type {mask_type!r}.") + normalized_mask_types[mask_type] = sequence_ids + attn_mask_type_per_seq = normalized_mask_types + + if window_size_per_mask_type is not None: + normalized_window_sizes = {} + for mask_type, policy_window_size in window_size_per_mask_type.items(): + mask_type = mask_type.replace(",", "_") + if mask_type == "causal_padding": + mask_type = "padding_causal" + if mask_type not in attn_mask_type_per_seq: + raise ValueError( + "window_size_per_mask_type contains a mask not present in " + f"attn_mask_type_per_seq: {mask_type!r}." + ) + normalized_window_sizes[mask_type] = policy_window_size + window_size_per_mask_type = normalized_window_sizes query_layer, key_layer, value_layer, declared_qkv_layout = _unpack_packed_qkv( qkv_layer, @@ -1909,7 +1861,11 @@ def forward( ) # checks for attention mask - if attn_mask_type is None: + if attn_mask_type_per_seq is not None: + # This scalar is used only for shared shape/window validation. + # Each policy call below uses its dictionary key. + attn_mask_type = next(iter(attn_mask_type_per_seq)) + elif attn_mask_type is None: attn_mask_type = self.attn_mask_type else: attn_mask_type = attn_mask_type.replace(",", "_") @@ -1986,6 +1942,84 @@ def forward( seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] max_seqlen_kv = int((seqlens_kv.max().item() + 63) // 64 * 64) + if attn_mask_type_per_seq is not None: + if qkv_format != "thd": + raise ValueError( + "attn_mask_type_per_seq is supported only with qkv_format='thd'." + ) + if attention_mask is not None: + raise ValueError( + "attn_mask_type_per_seq does not support an explicit attention_mask." + ) + if ( + core_attention_bias_type != "no_bias" + or core_attention_bias is not None + or alibi_slopes is not None + ): + raise ValueError( + "attn_mask_type_per_seq does not yet support attention bias or ALiBi." + ) + if inference_params is not None: + raise ValueError( + "attn_mask_type_per_seq does not support KV caching." + ) + if ( + score_mod is not None + or score_mod_bprop is not None + or score_mod_tensors is not None + or score_mod_bprop_tensors is not None + ): + raise ValueError( + "attn_mask_type_per_seq does not support score modification." + ) + if fp8_output or self.fp8 or self.return_max_logit: + raise ValueError( + "attn_mask_type_per_seq does not yet support FP8 or max-logit output." + ) + assigned_sequence_count = 0 + for sequence_ids in attn_mask_type_per_seq.values(): + if not isinstance(sequence_ids, torch.Tensor): + raise ValueError( + "attn_mask_type_per_seq values must be sequence-ID tensors." + ) + if sequence_ids.ndim != 1 or sequence_ids.dtype not in ( + torch.int32, + torch.int64, + ): + raise ValueError( + "attn_mask_type_per_seq values must be one-dimensional int32 or " + "int64 sequence-ID tensors." + ) + if not sequence_ids.is_cuda or sequence_ids.device != query_layer.device: + raise ValueError( + "attn_mask_type_per_seq sequence IDs must be on the Q/K/V device." + ) + assigned_sequence_count += sequence_ids.numel() + if assigned_sequence_count != batch_size: + raise ValueError( + "attn_mask_type_per_seq must assign every sequence exactly once by count: " + f"got {assigned_sequence_count} assignments for {batch_size} sequences." + ) + return self._forward_thd_mask_types( + query_layer, + key_layer, + value_layer, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + attn_mask_type_per_seq, + window_size_per_mask_type, + max_seqlen_q, + max_seqlen_kv, + window_size=window_size, + checkpoint_core_attention=checkpoint_core_attention, + fast_zero_fill=fast_zero_fill, + pad_between_seqs=pad_between_seqs, + bf16_backward=bf16_backward, + num_splits=num_splits, + ) + # update KV cache and retrieve saved tokens from cache for inference if inference_params is not None: assert self.layer_number is not None, "Layer number must be set!" From 98dcf173445c421171c01fa000c34cd0198ac5d1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 18:00:10 +0000 Subject: [PATCH 4/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../attention/test_mixed_thd_attention.py | 8 ++------ .../dot_product_attention.py | 16 ++++------------ 2 files changed, 6 insertions(+), 18 deletions(-) diff --git a/tests/pytorch/attention/test_mixed_thd_attention.py b/tests/pytorch/attention/test_mixed_thd_attention.py index 4cb3f39578..5d6fd6f16f 100644 --- a/tests/pytorch/attention/test_mixed_thd_attention.py +++ b/tests/pytorch/attention/test_mixed_thd_attention.py @@ -62,9 +62,7 @@ def _per_sequence_scalar_reference( for mask_type, sequence_ids in mask_type_per_seq.items(): for sequence_id in sequence_ids.tolist(): q_length = int((cu_seqlens_q[sequence_id + 1] - cu_seqlens_q[sequence_id]).item()) - kv_length = int( - (cu_seqlens_kv[sequence_id + 1] - cu_seqlens_kv[sequence_id]).item() - ) + kv_length = int((cu_seqlens_kv[sequence_id + 1] - cu_seqlens_kv[sequence_id]).item()) q_start = int(q_physical[sequence_id].item()) kv_start = int(kv_physical[sequence_id].item()) q_end = q_start + q_length @@ -231,9 +229,7 @@ def test_thd_mask_types_support_cross_attention_and_per_policy_windows(): dtype = torch.float16 cu_seqlens_q = _make_cu_seqlens((5, 3, 4)) cu_seqlens_kv = _make_cu_seqlens((7, 2, 5)) - mask_type_per_seq = _make_policy( - {"padding": (0, 2), "padding_causal_bottom_right": (1,)} - ) + mask_type_per_seq = _make_policy({"padding": (0, 2), "padding_causal_bottom_right": (1,)}) windows = {"padding": (-1, -1), "padding_causal_bottom_right": (2, 0)} query = ( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 4d286852cb..e3a177b047 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1342,9 +1342,7 @@ def _select_thd_sequences( torch.cumsum(sequence_lengths, dim=0, dtype=torch.int32), ) ) - selected_cu_seqlens_padded = torch.cat( - (physical_starts, physical_cu_seqlens[-1:]) - ) + selected_cu_seqlens_padded = torch.cat((physical_starts, physical_cu_seqlens[-1:])) return selected_cu_seqlens, selected_cu_seqlens_padded def _forward_thd_mask_types( @@ -1746,9 +1744,7 @@ def forward( """ if window_size_per_mask_type is not None and attn_mask_type_per_seq is None: - raise ValueError( - "window_size_per_mask_type requires attn_mask_type_per_seq." - ) + raise ValueError("window_size_per_mask_type requires attn_mask_type_per_seq.") if attn_mask_type_per_seq is not None: if not isinstance(attn_mask_type_per_seq, dict) or not attn_mask_type_per_seq: raise ValueError("attn_mask_type_per_seq must be a non-empty dictionary.") @@ -1960,18 +1956,14 @@ def forward( "attn_mask_type_per_seq does not yet support attention bias or ALiBi." ) if inference_params is not None: - raise ValueError( - "attn_mask_type_per_seq does not support KV caching." - ) + raise ValueError("attn_mask_type_per_seq does not support KV caching.") if ( score_mod is not None or score_mod_bprop is not None or score_mod_tensors is not None or score_mod_bprop_tensors is not None ): - raise ValueError( - "attn_mask_type_per_seq does not support score modification." - ) + raise ValueError("attn_mask_type_per_seq does not support score modification.") if fp8_output or self.fp8 or self.return_max_logit: raise ValueError( "attn_mask_type_per_seq does not yet support FP8 or max-logit output." From d52a461478ce4ffab5ba9b41a6cb3fea9e9898f1 Mon Sep 17 00:00:00 2001 From: Desh Raj Date: Fri, 7 Aug 2026 14:20:49 -0700 Subject: [PATCH 5/5] Mask inactive mixed-THD lanes for FA3 Mixed-mask dispatch represents sequences owned by other policies as inter-sequence padding. FA3 preserves that layout contract, but its padding lanes are unspecified: focused Hopper testing observed NaNs in inactive forward outputs and tile-spill garbage in inactive dQ/dK/dV lanes. Build a sync-free mask from the physical cu_seqlens, sanitize each policy output with torch.where, and wrap Q/K/V in an identity-forward autograd guard that zeros only inactive gradient lanes. This keeps the pad_between_seqs design, avoids regrouping or copying tokens, and lets disjoint policy results be summed safely across FA2 and FA3. Validation: NRT H100 source-overlay run passed all 16 focused mixed-THD tests for FP16/BF16 forward and backward, pre-existing padding, cross attention, per-policy windows, CUDA graphs, and validation errors. Signed-off-by: Desh Raj --- .../dot_product_attention.py | 71 +++++++++++++++++-- 1 file changed, 66 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index e3a177b047..c3a0a509e5 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -87,6 +87,20 @@ } +class _IdentityWithMaskedGradient(torch.autograd.Function): + """Preserve tensor storage in forward and zero unselected backward lanes.""" + + @staticmethod + def forward(ctx, tensor: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + ctx.save_for_backward(mask) + return tensor + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]: + (mask,) = ctx.saved_tensors + return torch.where(mask, grad_output, 0.0), None + + def _infer_custom_dpa_local_recipes( fp8_recipe: Recipe, fp8_meta: Dict[str, Any], @@ -1345,6 +1359,30 @@ def _select_thd_sequences( selected_cu_seqlens_padded = torch.cat((physical_starts, physical_cu_seqlens[-1:])) return selected_cu_seqlens, selected_cu_seqlens_padded + @staticmethod + def _selected_thd_token_mask( + token_count: int, + cu_seqlens: torch.Tensor, + cu_seqlens_padded: Optional[torch.Tensor], + sequence_ids: torch.Tensor, + ) -> torch.Tensor: + """Return a sync-free mask for selected tokens in physical THD storage.""" + sequence_ids = sequence_ids.to(dtype=torch.long) + physical_cu_seqlens = cu_seqlens if cu_seqlens_padded is None else cu_seqlens_padded + starts = physical_cu_seqlens.index_select(0, sequence_ids) + lengths = cu_seqlens.index_select(0, sequence_ids + 1) - cu_seqlens.index_select( + 0, sequence_ids + ) + ends = starts + lengths + boundaries = torch.cat((starts, ends)).to(dtype=torch.long) + deltas = torch.cat((torch.ones_like(starts), -torch.ones_like(ends))) + token_deltas = torch.zeros( + token_count + 1, + dtype=deltas.dtype, + device=cu_seqlens.device, + ).index_add(0, boundaries, deltas) + return torch.cumsum(token_deltas[:-1], dim=0, dtype=torch.int32).ne(0) + def _forward_thd_mask_types( self, query_layer: torch.Tensor, @@ -1370,8 +1408,9 @@ def _forward_thd_mask_types( Each scalar-mask call sees the original Q/K/V storage. Sequences assigned to other mask types are represented as inter-sequence padding, so the - backend leaves zeros at those positions. The disjoint outputs can then be - added without a token gather or scatter. + backend is allowed to leave unspecified values at those positions. Each + policy output is masked to its selected physical tokens before the + disjoint outputs are added, without a token gather or scatter. """ output_features = query_layer.shape[-2] * value_layer.shape[-1] if query_layer.shape[0] == 0: @@ -1415,6 +1454,27 @@ def _forward_thd_mask_types( output = None for mask_type in nonempty_mask_types: sequence_ids = attn_mask_type_per_seq[mask_type] + policy_q_token_mask = self._selected_thd_token_mask( + query_layer.shape[0], + cu_seqlens_q, + cu_seqlens_q_padded, + sequence_ids, + ) + policy_kv_token_mask = self._selected_thd_token_mask( + key_layer.shape[0], + cu_seqlens_kv, + cu_seqlens_kv_padded, + sequence_ids, + ) + policy_query = _IdentityWithMaskedGradient.apply( + query_layer, policy_q_token_mask[:, None, None] + ) + policy_key = _IdentityWithMaskedGradient.apply( + key_layer, policy_kv_token_mask[:, None, None] + ) + policy_value = _IdentityWithMaskedGradient.apply( + value_layer, policy_kv_token_mask[:, None, None] + ) selected_cu_seqlens_q, selected_cu_seqlens_q_padded = self._select_thd_sequences( cu_seqlens_q, cu_seqlens_q_padded, @@ -1429,9 +1489,9 @@ def _forward_thd_mask_types( if window_size_per_mask_type is not None: policy_window_size = window_size_per_mask_type.get(mask_type, window_size) policy_output = self.forward( - query_layer, - key_layer, - value_layer, + policy_query, + policy_key, + policy_value, qkv_format="thd", cu_seqlens_q=selected_cu_seqlens_q, cu_seqlens_kv=selected_cu_seqlens_kv, @@ -1447,6 +1507,7 @@ def _forward_thd_mask_types( bf16_backward=bf16_backward, num_splits=num_splits, ) + policy_output = torch.where(policy_q_token_mask.unsqueeze(-1), policy_output, 0.0) output = policy_output if output is None else output + policy_output if output is None: raise ValueError("attn_mask_type_per_seq must assign at least one sequence.")