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 new file mode 100644 index 0000000000..ea1d726dc8 --- /dev/null +++ b/tests/pytorch/attention/test_mixed_thd_attention.py @@ -0,0 +1,423 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for per-sequence mask types in packed THD attention.""" + +import pytest +import torch + +from transformer_engine.pytorch import DotProductAttention +from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, +) + + +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, 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_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 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_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, + ) + + 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( + "mask_type_per_seq", + ( + {"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_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) + 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 ( + 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) + output = 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, + ) + reference_output = _per_sequence_scalar_reference( + attention, + reference_query, + reference_key, + reference_value, + cu_seqlens, + 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(output, reference_output, **tolerances) + + 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, **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) + + +@pytest.mark.parametrize( + ("compute_capability", "fa3_installed", "flash_enabled", "fa3_enabled", "expected"), + ( + ((9, 0), True, "1", "1", True), + ((8, 0), True, "1", "1", False), + ((10, 0), True, "1", "1", False), + ((9, 0), False, "1", "1", False), + ((9, 0), True, "0", "1", False), + ((9, 0), True, "1", "0", False), + ), +) +def test_thd_mask_type_runtime_dispatch( + monkeypatch, + compute_capability, + fa3_installed, + flash_enabled, + fa3_enabled, + expected, +): + """Only Hopper with enabled FA3 uses the inter-sequence-padding implementation.""" + monkeypatch.setattr( + dpa_module.dpa_utils, "get_device_compute_capability", lambda: compute_capability + ) + monkeypatch.setattr(dpa_module.FlashAttentionUtils, "v3_is_installed", fa3_installed) + monkeypatch.setenv("NVTE_FLASH_ATTN", flash_enabled) + monkeypatch.setenv("NVTE_FLASH_ATTN_V3", fa3_enabled) + + assert dpa_module.dpa_utils.is_thd_mask_type_padding_supported() is expected + + +def test_thd_mask_types_support_cuda_graph_capture(): + """The sync-free policy metadata path must be replayable in a CUDA graph.""" + if not dpa_module.dpa_utils.is_thd_mask_type_padding_supported(): + pytest.skip("CUDA graph capture requires the Hopper/FA3 padding implementation") + + 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 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_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 + + if invalid_case == "policy_dtype": + 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: + windows = {"padding_causal_bottom_right": (-1, 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, + attn_mask_type_per_seq=policies, + window_size_per_mask_type=windows, + ) + + +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 + ) + 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, + 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() + 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..ad5725c8af 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], @@ -1323,6 +1337,314 @@ def get_quantizer_roles( ] return base[:num_quantizers] + @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 + + @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_with_padding( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: 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, + bf16_backward: Optional[bool], + num_splits: Optional[int], + ) -> torch.Tensor: + """Dispatch mixed THD masks with FA3 inter-sequence padding.""" + nonempty_mask_types = [ + mask_type + for mask_type, sequence_ids in attn_mask_type_per_seq.items() + if sequence_ids.numel() > 0 + ] + 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, + 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( + policy_query, + policy_key, + policy_value, + 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, + ) + 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.") + return output + + def _forward_thd_mask_types_grouped( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: 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, + bf16_backward: Optional[bool], + num_splits: Optional[int], + ) -> torch.Tensor: + """Dispatch mixed THD masks by compacting each policy for FA2.""" + output_features = query_layer.shape[-2] * value_layer.shape[-1] + output = query_layer.new_zeros((query_layer.shape[0], output_features)) + for mask_type, sequence_ids in attn_mask_type_per_seq.items(): + if sequence_ids.numel() == 0: + continue + q_token_mask = self._selected_thd_token_mask( + query_layer.shape[0], + cu_seqlens_q, + cu_seqlens_q_padded, + sequence_ids, + ) + kv_token_mask = self._selected_thd_token_mask( + key_layer.shape[0], + cu_seqlens_kv, + cu_seqlens_kv_padded, + sequence_ids, + ) + q_token_indices = torch.nonzero(q_token_mask, as_tuple=False).flatten() + kv_token_indices = torch.nonzero(kv_token_mask, as_tuple=False).flatten() + selected_cu_seqlens_q, _ = self._select_thd_sequences( + cu_seqlens_q, + cu_seqlens_q_padded, + sequence_ids, + ) + selected_cu_seqlens_kv, _ = 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.index_select(0, q_token_indices), + key_layer.index_select(0, kv_token_indices), + value_layer.index_select(0, kv_token_indices), + qkv_format="thd", + cu_seqlens_q=selected_cu_seqlens_q, + cu_seqlens_kv=selected_cu_seqlens_kv, + 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=False, + bf16_backward=bf16_backward, + num_splits=num_splits, + ) + output = output.index_copy(0, q_token_indices, policy_output) + return output + + def _forward_thd_mask_types( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: 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 THD masks through the best runtime implementation.""" + output_features = query_layer.shape[-2] * value_layer.shape[-1] + 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()) + + 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( + query_layer, + key_layer, + value_layer, + qkv_format="thd", + 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, + ) + if not nonempty_mask_types: + raise ValueError("attn_mask_type_per_seq must assign at least one sequence.") + + if dpa_utils.is_thd_mask_type_padding_supported(): + return self._forward_thd_mask_types_with_padding( + 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, + bf16_backward=bf16_backward, + num_splits=num_splits, + ) + return self._forward_thd_mask_types_grouped( + 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, + bf16_backward=bf16_backward, + num_splits=num_splits, + ) + @no_torch_dynamo(recursive=False) def forward( self, @@ -1357,6 +1679,8 @@ def forward( qkv_layer: Optional[torch.Tensor] = None, kv_layer: Optional[torch.Tensor] = None, qkv_interleave_dim: int = -3, + 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. @@ -1596,8 +1920,60 @@ 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. + 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. On Hopper with FlashAttention 3 enabled, Transformer Engine + keeps Q/K/V in their original packed layout and represents sequences owned + by other policies as inter-sequence padding. Otherwise it compacts each + policy's logical tokens for an ordinary scalar-mask call and restores the + original physical THD layout. The caller-facing API is identical for both + implementations. + 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 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, kv_layer, @@ -1676,7 +2052,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(",", "_") @@ -1753,6 +2133,80 @@ 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!" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 6eb3ce54f1..39d7a7e376 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -379,6 +379,22 @@ def _get_fused_attn_backend( ) +def is_thd_mask_type_padding_supported() -> bool: + """Whether mixed THD masks can use the inter-sequence-padding implementation. + + FlashAttention 3 is the FlashAttention backend that supports padding between + sequences, and TE only enables it on Hopper (SM90). Keep this capability + selection next to the rest of the attention backend architecture filters so + callers can fall back without exposing a hardware-specific API. + """ + return ( + get_device_compute_capability() == (9, 0) + and FlashAttentionUtils.v3_is_installed + and bool(int(os.environ.get("NVTE_FLASH_ATTN", "1"))) + and bool(int(os.environ.get("NVTE_FLASH_ATTN_V3", "1"))) + ) + + def get_attention_backend( attention_params: AttentionParams = None, ):