diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index da4f024dc72b..a3e37b5176e7 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -174,7 +174,7 @@ Tile geometry and blending follow the loaded Diffusers VAE. Tensor parallelism shards the DiT attention and gated MLP weights via `parallel_config.tp_size`. The text encoder, AdaLN modulation, and media input/output projections remain replicated. TP can be combined with Ulysses; -Ulysses requires VANILLA or FA4 attention. CFG, ring, and Attention2D parallelism +Ulysses supports VANILLA, FA4, TRTLLM, and FLASHINFER attention. CFG, ring, and Attention2D parallelism are not supported. See the [2-GPU TP config](../../../examples/visual_gen/configs/minimax-h3-bf16-tp2.yaml). diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py index c660168c64e6..78bd4690e424 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py @@ -444,6 +444,79 @@ def _compact_qkv( v = v.reshape(batch_size * kv_seq_len, -1).contiguous() return q, k, v + @torch.compiler.disable + def _forward_padded_self_attention( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + key_padding_mask: torch.Tensor, + timestep: Optional[torch.Tensor], + ) -> torch.Tensor: + """Attend to a valid prefix; callers explicitly discard padded query rows. + + Ulysses still needs the padded output shape for its return collective. + Cache prefix lengths across layers sharing the same request mask to avoid + repeating the device-to-host synchronization in every attention block. + """ + if key_padding_mask.dtype != torch.bool or key_padding_mask.shape != q.shape[:2]: + raise ValueError("TRTLLM padding mask must be boolean with shape [batch, sequence].") + state = self.metadata._metadata_state + cached = state.get("padding_prefix_lengths") + version = None if torch.is_inference(key_padding_mask) else key_padding_mask._version + if ( + version is None + or cached is None + or cached[0] is not key_padding_mask + or cached[1] != version + ): + mask = key_padding_mask.cpu() + lengths = mask.sum(dim=1).tolist() + expected = torch.arange(mask.shape[1]).unsqueeze(0) < torch.tensor(lengths).unsqueeze(1) + if not torch.equal(mask, expected): + raise ValueError("TRTLLM padded self-attention requires trailing padding.") + state["padding_prefix_lengths"] = (key_padding_mask, version, lengths) + else: + lengths = cached[2] + batch_size, seq_len = q.shape[:2] + if all(length == seq_len for length in lengths): + return self.forward(q, k, v, batch_size, seq_len, timestep=timestep) + + groups: dict[int, list[int]] = {} + for batch_idx, length in enumerate(lengths): + groups.setdefault(length, []).append(batch_idx) + if len(groups) == 1: + length = lengths[0] + if length == 0: + return q.new_zeros((batch_size, seq_len, self.num_heads * self.head_dim)) + output = self.forward( + q=q[:, :length].contiguous(), + k=k[:, :length].contiguous() if k is not None else None, + v=v[:, :length].contiguous() if v is not None else None, + batch_size=batch_size, + seq_len=length, + timestep=timestep, + ) + return torch.nn.functional.pad(output, (0, 0, 0, seq_len - length)) + + output = q.new_zeros((batch_size, seq_len, self.num_heads * self.head_dim)) + for length, batch_indices in groups.items(): + if length == 0: + continue + indices = torch.tensor(batch_indices, device=q.device) + group_output = self.forward( + q=q[:, :length].index_select(0, indices), + k=k[:, :length].index_select(0, indices) if k is not None else None, + v=v[:, :length].index_select(0, indices) if v is not None else None, + batch_size=len(batch_indices), + seq_len=length, + timestep=timestep, + ) + output.index_copy_( + 0, indices, torch.nn.functional.pad(group_output, (0, 0, 0, seq_len - length)) + ) + return output + def forward( self, q: torch.Tensor, @@ -456,6 +529,8 @@ def forward( sparse_backend_args: Optional[SparseBackendForwardArgs] = None, timestep: Optional[torch.Tensor] = None, kv: Optional[ForwardPlan] = None, + key_padding_mask: Optional[torch.Tensor] = None, + discard_padded_queries: bool = False, **kwargs, ) -> torch.Tensor: """ @@ -485,11 +560,6 @@ def forward( prediction hooks. A ``block_sparse_inputs`` payload selects the generic block-sparse FMHA. timestep: Denoising timestep consumed by the sparse prediction hooks. - **kwargs: Backend-specific keyword arguments forwarded by the attention - module, such as ``key_padding_mask``; ignored here, as by the other - backends. The TRTLLM kernels have no key-padding input, so callers - that need padded keys must guard at the model level or select the - ``VANILLA`` backend. kv: A ``ForwardPlan`` from a ``BlockCausalKVCacheManager``. When given, ``k``/``v`` are the new tokens only: the fused kernel writes them at ``staging_offset`` and attends over everything cached before them plus @@ -497,10 +567,31 @@ def forward( The first ``kv.num_tokens`` rows are real; rows past them are padding, neither cached nor attended, and output zero. ``batch_size`` must be 1. Sparse attention is not supported on this path. + key_padding_mask: Boolean valid-prefix mask, honored only when + discard_padded_queries=True. Otherwise retains the legacy TRTLLM + behavior of ignoring this backend-specific keyword. + discard_padded_queries: Opt in to valid-prefix dense self-attention. + Caller guarantees masked query rows are unused; their outputs + are zero-filled. Callers requiring general masking must select + a backend that supports it, such as VANILLA. + **kwargs: Additional backend-specific arguments. Returns: Output tensor [B, S, H*D] """ + if key_padding_mask is not None and discard_padded_queries: + if ( + kv is not None + or attention_mask != PredefinedAttentionMask.FULL + or seq_len_kv not in (None, seq_len) + or self.quant_attention_config is not None + or sparse_backend_args is not None + or self.sparse_params is not None + ): + raise ValueError( + "TRTLLM padding requires dense self-attention with discarded padded queries." + ) + return self._forward_padded_self_attention(q, k, v, key_padding_mask, timestep) if kv is not None: if sparse_backend_args is not None: raise NotImplementedError("Sparse attention is not supported over a K/V cache.") diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py index 76751febecd2..2cc3999591eb 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py @@ -229,9 +229,14 @@ def __init__(self, pipeline_config: DiffusionPipelineConfig) -> None: "cfg_size=1, ring_size=1, attn2d_size=(1, 1), and " "tp_size * ulysses_size=world_size." ) - if vgm.ulysses_size > 1 and pipeline_config.attention.backend not in ("VANILLA", "FA4"): + if vgm.ulysses_size > 1 and pipeline_config.attention.backend not in ( + "VANILLA", + "FA4", + "TRTLLM", + "FLASHINFER", + ): raise NotImplementedError( - "MiniMax-H3 Ulysses requires VANILLA or FA4 attention. " + "MiniMax-H3 Ulysses requires VANILLA, FA4, TRTLLM, or FLASHINFER attention. " "Other attention backends do not support the required " "packed-row padding mask." ) diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py index e09d334f917a..d748960e9e07 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/transformer_minimax_h3.py @@ -245,6 +245,7 @@ def forward( key, value, key_padding_mask=key_padding_mask, + discard_padded_queries=True, timestep=timestep, ) return self.to_out[0](hidden_states) @@ -285,6 +286,7 @@ def compute_v() -> torch.Tensor: compute_k, compute_v, key_padding_mask=key_padding_mask, + discard_padded_queries=True, timestep=timestep, ) batch_size, local_seq_len = out_4d.shape[:2] @@ -477,6 +479,8 @@ def __init__( self._supports_key_padding_mask = model_config.attention.backend.upper() in ( "VANILLA", "FA4", + "TRTLLM", + "FLASHINFER", ) self.sharder = SequenceSharder.from_vgm( model_config.visual_gen_mapping, @@ -512,10 +516,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if pad: if not self._supports_key_padding_mask: raise NotImplementedError( - "Padded token-refiner sequences require key_padding_mask support; use VANILLA or FA4." + "Padded token-refiner sequences require key_padding_mask support; " + "use VANILLA, FA4, TRTLLM or FLASHINFER." ) - valid = hidden_states.new_ones((hidden_states.shape[0], seq_len), dtype=torch.bool) - key_padding_mask = _pad_tensor_dim(valid, 1, pad, False) + # Backends cache mask plans using the tensor mutation counter. + with torch.inference_mode(False): + valid = hidden_states.new_ones((hidden_states.shape[0], seq_len), dtype=torch.bool) + key_padding_mask = _pad_tensor_dim(valid, 1, pad, False) hidden_states = _pad_tensor_dim(hidden_states, 1, pad) hidden_states = self.sharder.shard(hidden_states, dim=1) for block in self.refiner_blocks: @@ -666,7 +673,12 @@ def __init__(self, model_config: DiffusionModelConfig) -> None: super().__init__(model_config) self._attn_backend = getattr(model_config.attention, "backend", "VANILLA") # Only VANILLA supports arbitrary padding between packed modalities. - self._supports_key_padding_mask = self._attn_backend in ("VANILLA", "FA4") + self._supports_key_padding_mask = self._attn_backend in ( + "VANILLA", + "FA4", + "TRTLLM", + "FLASHINFER", + ) quant_algo = model_config.quant_config.quant_algo # TODO: Support calibrated checkpoints in load_weights before enabling static quantization. if quant_algo is not None and not ( @@ -997,19 +1009,23 @@ def forward( "Padded packed sequences (negative token_tags) need an attention " "backend that honours key_padding_mask; the " f"{self._attn_backend} backend cannot express it. Use " - "VANILLA or FA4." + "VANILLA, FA4, TRTLLM or FLASHINFER." ) - # FA4 expresses padding through seqused_k, which requires valid - # tokens to form a prefix. Ulysses appends padding in this layout. - if self._attn_backend == "FA4" and bool( + # FA4 seqused_k and TRTLLM prefix trimming require trailing padding. + # Ulysses appends padding in this layout. + if self._attn_backend in ("FA4", "TRTLLM") and bool( ((token_tags_for_mask[:-1] < 0) & (token_tags_for_mask[1:] >= 0)).any() ): raise NotImplementedError( - "FA4 requires trailing packed-row padding; use VANILLA for interior padding." + f"{self._attn_backend} requires trailing packed-row padding; use VANILLA for interior padding." + ) + # Preserve a mutation counter for shared backend mask-plan caches. + with torch.inference_mode(False): + key_padding_mask = ( + (token_tags_for_mask >= 0) + .unsqueeze(0) + .expand(packed_hidden_states.shape[0], -1) ) - key_padding_mask = ( - (token_tags_for_mask >= 0).unsqueeze(0).expand(packed_hidden_states.shape[0], -1) - ) padded_sequence_length = packed_hidden_states.shape[1] packed_hidden_states = self.sharder.shard(packed_hidden_states, dim=1) diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 06bd8dfc616a..8fbfc095ee03 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -506,7 +506,8 @@ def test_minimax_h3_text_encoder_tp_lpips(tmp_path, monkeypatch, world_size): @pytest.mark.parametrize( - "ulysses_size,vae_size,backend", [(2, 2, "VANILLA"), (4, 3, "VANILLA"), (2, 2, "FA4")] + "ulysses_size,vae_size,backend", + [(2, 2, "VANILLA"), (4, 3, "VANILLA"), (2, 2, "FA4"), (2, 2, "TRTLLM"), (2, 2, "FLASHINFER")], ) def test_minimax_h3_ulysses_parallel_vae_lpips(tmp_path, ulysses_size, vae_size, backend): """Compare real Ulysses + tiled VAE collectives with a fresh one-GPU baseline.""" diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 5dfdc6fbcdd4..0ca0c33f142c 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -41,12 +41,16 @@ l0_dgx_b200: - unittest/_torch/visual_gen/multi_gpu/test_ltx2_parallel_vae.py - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[VANILLA-2-1] - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FA4-2-1] + - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[TRTLLM-2-1] + - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FLASHINFER-2-1] - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_ulysses_and_parallel_vae - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_tp2_example_lpips TIMEOUT (30) - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_text_encoder_tp.py::test_minimax_h3_text_encoder_tp[tp2_world2] - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_text_encoder_tp_lpips[2] TIMEOUT (30) - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[2-2-VANILLA] TIMEOUT (30) - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[2-2-FA4] TIMEOUT (30) + - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[2-2-TRTLLM] TIMEOUT (30) + - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[2-2-FLASHINFER] TIMEOUT (30) - unittest/_torch/visual_gen/test_flux_pipeline.py::TestFluxParallelism::test_ulysses_2gpu_correctness - unittest/_torch/visual_gen/test_flux_pipeline.py::TestFluxCombinedOptimizations::test_all_optimizations_combined - condition: @@ -67,8 +71,12 @@ l0_dgx_b200: - unittest/disaggregated/test_transfer_lifecycle_gpu.py - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[VANILLA-4-1] - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FA4-4-1] + - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[TRTLLM-4-1] + - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FLASHINFER-4-1] - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[VANILLA-2-2] - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FA4-2-2] + - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[TRTLLM-2-2] + - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FLASHINFER-2-2] - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_partial_vae_groups - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_text_encoder_tp.py::test_minimax_h3_text_encoder_tp[tp2_world4] - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_text_encoder_tp_lpips[4] TIMEOUT (30) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py index d8a01c9d3de7..a581c9f58274 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py @@ -38,8 +38,8 @@ def _make_model_config( return DiffusionModelConfig( pretrained_config=SimpleNamespace( num_attention_heads=2, - attention_head_dim=128 if backend == "FA4" else 16, - hidden_size=256 if backend == "FA4" else 32, + attention_head_dim=128 if backend in ("FA4", "TRTLLM", "FLASHINFER") else 16, + hidden_size=256 if backend in ("FA4", "TRTLLM", "FLASHINFER") else 32, num_layers=2, num_refiner_layers=1, ffn_dim=32, @@ -190,8 +190,10 @@ def _logic_h3_ulysses_and_parallel_vae( @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires two CUDA GPUs") @pytest.mark.parametrize("async_ulysses", [False, True], ids=["nccl", "async"]) -@pytest.mark.parametrize("backend", ["VANILLA", "FA4"]) +@pytest.mark.parametrize("backend", ["VANILLA", "FA4", "TRTLLM", "FLASHINFER"]) def test_h3_ulysses_and_parallel_vae(backend: str, async_ulysses: bool) -> None: + if backend == "TRTLLM" and torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("TRTLLM H3 requires SM100 or SM103") if backend == "FA4" and torch.cuda.get_device_capability()[0] < 10: pytest.skip("FA4 requires SM100 or newer") if torch.cuda.device_count() < 2: @@ -261,11 +263,13 @@ def test_h3_partial_vae_groups() -> None: @pytest.mark.parametrize("tp_size,ulysses_size", [(2, 1), (4, 1), (2, 2)]) -@pytest.mark.parametrize("backend", ["VANILLA", "FA4"]) +@pytest.mark.parametrize("backend", ["VANILLA", "FA4", "TRTLLM", "FLASHINFER"]) def test_h3_tensor_parallel(tp_size: int, ulysses_size: int, backend: str) -> None: world_size = tp_size * ulysses_size if torch.cuda.device_count() < world_size: pytest.skip(f"Requires {world_size} CUDA GPUs") + if backend == "TRTLLM" and torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("TRTLLM H3 requires SM100 or SM103") if backend == "FA4" and torch.cuda.get_device_capability()[0] < 10: pytest.skip("FA4 requires SM100 or newer") run_distributed( diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py index 92f229135c1c..aed9546cfc8a 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py @@ -1649,7 +1649,7 @@ def test_pipeline_rejects_unsupported_ulysses_topology(field: str, value: int) - MiniMaxH3Pipeline(config) -@pytest.mark.parametrize("backend", ["TRTLLM", "CUTEDSL"]) +@pytest.mark.parametrize("backend", ["CUTEDSL", "CUDNN"]) def test_pipeline_rejects_unvalidated_ulysses_attention(backend: str) -> None: config = _ulysses_pipeline_config() config.attention.backend = backend @@ -1780,3 +1780,16 @@ def initialize(self, config): config.visual_gen_mapping.ulysses_size = ulysses_size MiniMaxH3Pipeline(config) assert initialized == [config] + + +@pytest.mark.parametrize("backend", ["TRTLLM", "FLASHINFER"]) +def test_pipeline_accepts_masked_ulysses_backends( + monkeypatch: pytest.MonkeyPatch, backend: str +) -> None: + monkeypatch.setattr( + h3_pipeline.BasePipeline, "__init__", lambda self, config: torch.nn.Module.__init__(self) + ) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (10, 0)) + config = _ulysses_pipeline_config(2) + config.attention.backend = backend + MiniMaxH3Pipeline(config) diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py index 2ed43d61e447..aae16f0e16a0 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_transformer.py @@ -935,7 +935,7 @@ def _cpu_swiglu(hidden_states: torch.Tensor, **kwargs: object) -> torch.Tensor: @pytest.mark.parametrize( ("backend", "supported"), - [("VANILLA", True), ("FA4", True), ("TRTLLM", False), ("CUTEDSL", False)], + [("VANILLA", True), ("FA4", True), ("TRTLLM", True), ("FLASHINFER", True), ("CUTEDSL", False)], ) def test_key_padding_mask_support_tracks_attention_backend(backend: str, supported: bool) -> None: config = _make_model_config() @@ -947,7 +947,7 @@ def test_key_padding_mask_support_tracks_attention_backend(backend: str, support @requires_cuda def test_forward_rejects_padded_rows_on_backend_that_ignores_the_mask() -> None: - """TRTLLM/CUTEDSL swallow key_padding_mask via **kwargs, so fail loudly instead.""" + """Reject padding when the selected backend cannot express it.""" model = h3.MiniMaxH3Transformer3DModel(_make_model_config()).to("cuda") _initialize_weights(model) inputs = _model_inputs("cuda") @@ -958,7 +958,7 @@ def test_forward_rejects_padded_rows_on_backend_that_ignores_the_mask() -> None: model(**inputs) model._supports_key_padding_mask = False - model._attn_backend = "TRTLLM" + model._attn_backend = "CUTEDSL" with pytest.raises(NotImplementedError, match="key_padding_mask"): model(**inputs) diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py index a8dad92f8c06..8d9044ac6e0f 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -7,6 +7,7 @@ import pytest import torch +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask from tensorrt_llm._torch.attention.backends.sparse.params import ( BlockSparseForwardInputs, SparseBackendForwardArgs, @@ -376,3 +377,92 @@ def forward() -> torch.Tensor: torch.cuda.synchronize() torch.testing.assert_close(captured, eager) + + +@_REQUIRES_SM100 +@pytest.mark.parametrize("fused", [False, True]) +@pytest.mark.parametrize("lengths", [[5, 7], [5, 5], [8, 8], [5, 7, 5, 0], [0, 0]]) +def test_trtllm_padded_self_attention_matches_valid_rows(fused: bool, lengths: list[int]) -> None: + """Padding never contributes keys, including when ranks need unequal real lengths.""" + attention = visual_trtllm.TrtllmAttention( + num_heads=2, + head_dim=128, + dtype=torch.bfloat16, + attention_metadata_state=create_attention_metadata_state(), + ) + torch.manual_seed(42) + batch_size = len(lengths) + q, k, v = [ + torch.randn(batch_size, 8, 2, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3) + ] + mask = torch.arange(8, device="cuda")[None] < torch.tensor(lengths, device="cuda")[:, None] + expected = ( + torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=mask[:, None, None] + ) + .transpose(1, 2) + .flatten(2) + ) + query = torch.cat((q, k, v), dim=2) if fused else q + kwargs = dict( + batch_size=batch_size, seq_len=8, key_padding_mask=mask, discard_padded_queries=True + ) + forward = attention.forward + attention.forward = Mock(wraps=forward) + actual = forward(query, None if fused else k, None if fused else v, **kwargs) + assert attention.forward.call_count == len(set(lengths) - {0}) + attention.forward = forward + torch.testing.assert_close(actual[mask], expected[mask], atol=2e-2, rtol=2e-2) + assert torch.count_nonzero(actual[~mask]) == 0 + with pytest.raises(ValueError, match="discarded padded queries"): + attention.forward(q, k, v, kv=Mock(), **kwargs) + with pytest.raises(ValueError, match="discarded padded queries"): + attention.forward(q, k, v, attention_mask=PredefinedAttentionMask.CAUSAL, **kwargs) + mask[0] = True + mask[0, 1] = False + with pytest.raises(ValueError, match="trailing padding"): + attention.forward(q, k, v, **kwargs) + + +@_REQUIRES_SM100 +def test_trtllm_padding_requires_explicit_opt_in() -> None: + """Existing callers retain dense behavior unless they opt into prefix masking.""" + attention = visual_trtllm.TrtllmAttention( + num_heads=2, + head_dim=128, + dtype=torch.bfloat16, + attention_metadata_state=create_attention_metadata_state(), + ) + q, k, v = [torch.randn(2, 8, 2, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + mask = torch.zeros(2, 8, device="cuda", dtype=torch.bool) + attention._forward_padded_self_attention = Mock(side_effect=AssertionError("not opted in")) + expected = attention.forward(q, k, v, batch_size=2, seq_len=8) + actual = attention.forward(q, k, v, batch_size=2, seq_len=8, key_padding_mask=mask) + torch.testing.assert_close(actual, expected) + attention._forward_padded_self_attention.assert_not_called() + + +@pytest.mark.parametrize("mode", ["dense", "quant", "sparse"]) +def test_legacy_mask_preserves_attention_dispatch(monkeypatch, mode: str) -> None: + captured = {} + _capture_core_forward(monkeypatch, captured) + attention = _make_wrapper( + quant_attention_config=visual_trtllm.QuantAttentionConfig() if mode == "quant" else None + ) + if mode == "sparse": + attention.sparse_params = SkipSoftmaxParams() + q = torch.randn(1, 4, 2, 8) + mask = torch.tensor([[True, False, True, False]]) + backend_args = ( + SparseBackendForwardArgs(block_sparse_inputs=_make_block_sparse_inputs()) + if mode == "sparse" + else None + ) + kwargs = dict(batch_size=1, seq_len=4, sparse_backend_args=backend_args) + expected = attention.forward(q, q, q, **kwargs) + expected_args = captured["forward_args"] + actual = attention.forward(q, q, q, key_padding_mask=mask, **kwargs) + torch.testing.assert_close(actual, expected) + assert captured["forward_args"] == expected_args + with pytest.raises(ValueError): + attention.forward(q, q, q, key_padding_mask=mask, discard_padded_queries=True, **kwargs)