From f4a7d9198ebd6d54c48c859292eb7587046dcc23 Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:15:28 +0000 Subject: [PATCH 1/4] [TRTLLM-16806][feat] Enable TRTLLM Ulysses attention for MiniMax-H3 Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- docs/source/models/visual-generation.md | 2 +- .../visual_gen/attention/backends/trtllm.py | 76 +++++++++++++++++-- .../models/minimax_h3/pipeline_minimax_h3.py | 9 ++- .../minimax_h3/transformer_minimax_h3.py | 40 +++++++--- .../visual_gen/test_visual_gen_multi_gpu.py | 3 +- .../test_lists/test-db/l0_dgx_b200.yml | 2 + .../multi_gpu/test_minimax_h3_parallel.py | 12 ++- .../visual_gen/test_minimax_h3_pipeline.py | 15 +++- .../visual_gen/test_minimax_h3_transformer.py | 6 +- .../test_trtllm_attention_metadata.py | 32 ++++++++ 10 files changed, 168 insertions(+), 29 deletions(-) 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..805bb6054929 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py @@ -444,6 +444,57 @@ 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] + outputs = [] + for batch_idx, length in enumerate(lengths): + if length == 0: + output = q.new_zeros((1, q.shape[1], self.num_heads * self.head_dim)) + else: + output = self.forward( + q=q[batch_idx : batch_idx + 1, :length], + k=k[batch_idx : batch_idx + 1, :length] if k is not None else None, + v=v[batch_idx : batch_idx + 1, :length] if v is not None else None, + batch_size=1, + seq_len=length, + timestep=timestep, + ) + output = torch.nn.functional.pad(output, (0, 0, 0, q.shape[1] - length)) + outputs.append(output) + return torch.cat(outputs, dim=0) + def forward( self, q: torch.Tensor, @@ -456,6 +507,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 +538,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 +545,28 @@ 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 for dense self-attention. + discard_padded_queries: Caller guarantees masked query rows are unused. + Required with key_padding_mask; padded output rows are zero-filled. + **kwargs: Additional backend-specific arguments. Returns: Output tensor [B, S, H*D] """ + if key_padding_mask is not None: + if ( + kv is not None + or not discard_padded_queries + 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..4748f848f408 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -47,6 +47,8 @@ l0_dgx_b200: - 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: 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..05e38feda0fb 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"]) 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..25a899126daa 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -376,3 +376,35 @@ def forward() -> torch.Tensor: torch.cuda.synchronize() torch.testing.assert_close(captured, eager) + + +@_REQUIRES_SM100 +@pytest.mark.parametrize("fused", [False, True]) +def test_trtllm_padded_self_attention_matches_valid_rows(fused: bool) -> 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) + q, k, v = [torch.randn(2, 8, 2, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + mask = torch.arange(8, device="cuda")[None] < torch.tensor([5, 7], 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=2, seq_len=8, key_padding_mask=mask, discard_padded_queries=True) + actual = attention.forward(query, None if fused else k, None if fused else v, **kwargs) + 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, batch_size=2, seq_len=8, key_padding_mask=mask) + mask[0, 1] = False + with pytest.raises(ValueError, match="trailing padding"): + attention.forward(q, k, v, **kwargs) From b35d267fe914561ffd7653926ba5bc4ed54081c8 Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Thu, 29 Oct 2026 15:05:16 +0000 Subject: [PATCH 2/4] [ci watcher] fix Hunyuan attention padding compatibility Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- .../transformer_hunyuan_video1_5.py | 23 ++++++++++- .../test_hunyuan_video1_5_transformer.py | 39 +++++++++++++++++++ .../test_trtllm_attention_metadata.py | 3 ++ 3 files changed, 63 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/transformer_hunyuan_video1_5.py b/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/transformer_hunyuan_video1_5.py index 0e3ea948fe3b..87500536631e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/transformer_hunyuan_video1_5.py +++ b/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/transformer_hunyuan_video1_5.py @@ -298,6 +298,21 @@ def __init__( bias: bool = True, model_config: Optional[DiffusionModelConfig] = None, ) -> None: + model_config = model_config or DiffusionModelConfig() + if model_config.attention.backend == "TRTLLM": + # The text refiner needs general key padding, not the valid-prefix + # and discarded-query contract of TRTLLM padded self-attention. + model_config = model_config.model_copy( + update={ + "attention": model_config.attention.model_copy( + update={ + "backend": "VANILLA", + "quant_attention_config": None, + "sparse_attention_config": None, + } + ) + } + ) super().__init__( hidden_size=dim, num_attention_heads=num_attention_heads, @@ -305,7 +320,7 @@ def __init__( qkv_mode=QKVMode.FUSE_QKV, qk_norm=False, bias=bias, - config=model_config or DiffusionModelConfig(), + config=model_config, ) def forward( @@ -646,7 +661,11 @@ def forward( attention_mask, (seq_len - attention_mask.shape[1], 0), value=True ).bool() - hidden_states = self._attn_impl(query, key, value, key_padding_mask=key_padding_mask) + # Context tokens are packed valid-first; padded context rows remain + # masked in every block and do not contribute to the video output. + hidden_states = self._attn_impl( + query, key, value, key_padding_mask=key_padding_mask, discard_padded_queries=True + ) # hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) diff --git a/tests/unittest/_torch/visual_gen/test_hunyuan_video1_5_transformer.py b/tests/unittest/_torch/visual_gen/test_hunyuan_video1_5_transformer.py index a3d3ddf3193c..de633f1979eb 100644 --- a/tests/unittest/_torch/visual_gen/test_hunyuan_video1_5_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_hunyuan_video1_5_transformer.py @@ -281,3 +281,42 @@ def test_hunyuan_video1_5_allclose_to_hf(self): if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("valid_tokens", [[True, True, False, False], [False, True, False, True]]) +def test_trtllm_refiner_respects_key_padding(valid_tokens: list[bool]) -> None: + from tensorrt_llm._torch.visual_gen.attention.backends.vanilla import VanillaAttention + from tensorrt_llm._torch.visual_gen.models.hunyuan_video1_5.transformer_hunyuan_video1_5 import ( + HunyuanVideo15RefinerAttention, + ) + + config = _create_model_config(CONFIG, backend="TRTLLM") + refiner = HunyuanVideo15RefinerAttention(256, 2, 128, model_config=config) + refiner = refiner.to(device="cuda", dtype=torch.bfloat16).eval() + assert isinstance(refiner.attn, VanillaAttention) + assert config.attention.backend == "TRTLLM" + + torch.manual_seed(42) + with torch.no_grad(): + for parameter in refiner.parameters(): + parameter.normal_(std=0.02) + hidden = torch.randn(1, 4, 256, device="cuda", dtype=torch.bfloat16) + mask = torch.tensor([valid_tokens], device="cuda") + q, k, v = refiner.get_qkv(hidden) + expected = ( + F.scaled_dot_product_attention( + q.view(1, 4, 2, 128).transpose(1, 2), + k.view(1, 4, 2, 128).transpose(1, 2), + v.view(1, 4, 2, 128).transpose(1, 2), + attn_mask=mask[:, None, None, :], + ) + .transpose(1, 2) + .reshape(1, 4, 256) + ) + expected = refiner.to_out[0](expected) + actual = refiner(hidden, mask) + torch.testing.assert_close(actual, expected) + hidden[~mask] += 10 + perturbed = refiner(hidden, mask) + torch.testing.assert_close(perturbed[mask], actual[mask]) 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 25a899126daa..2a518bc68119 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, @@ -405,6 +406,8 @@ def test_trtllm_padded_self_attention_matches_valid_rows(fused: bool) -> None: assert torch.count_nonzero(actual[~mask]) == 0 with pytest.raises(ValueError, match="discarded padded queries"): attention.forward(q, k, v, batch_size=2, seq_len=8, key_padding_mask=mask) + with pytest.raises(ValueError, match="discarded padded queries"): + attention.forward(q, k, v, attention_mask=PredefinedAttentionMask.CAUSAL, **kwargs) mask[0, 1] = False with pytest.raises(ValueError, match="trailing padding"): attention.forward(q, k, v, **kwargs) From 05dcc4d963a40924a7fa394e08e0175231ed4839 Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:45:10 +0000 Subject: [PATCH 3/4] Address review feedback on opt-in Ulysses padding Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- .../visual_gen/attention/backends/trtllm.py | 63 +++++++++++++------ .../transformer_hunyuan_video1_5.py | 23 +------ .../test_lists/test-db/l0_dgx_b200.yml | 6 ++ .../test_hunyuan_video1_5_transformer.py | 39 ------------ .../visual_gen/test_minimax_h3_pipeline.py | 2 +- .../test_trtllm_attention_metadata.py | 62 ++++++++++++++++-- 6 files changed, 109 insertions(+), 86 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py index 805bb6054929..78bd4690e424 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py @@ -478,22 +478,44 @@ def _forward_padded_self_attention( state["padding_prefix_lengths"] = (key_padding_mask, version, lengths) else: lengths = cached[2] - outputs = [] + 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: - output = q.new_zeros((1, q.shape[1], self.num_heads * self.head_dim)) - else: - output = self.forward( - q=q[batch_idx : batch_idx + 1, :length], - k=k[batch_idx : batch_idx + 1, :length] if k is not None else None, - v=v[batch_idx : batch_idx + 1, :length] if v is not None else None, - batch_size=1, - seq_len=length, - timestep=timestep, - ) - output = torch.nn.functional.pad(output, (0, 0, 0, q.shape[1] - length)) - outputs.append(output) - return torch.cat(outputs, dim=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, @@ -545,18 +567,21 @@ 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 for dense self-attention. - discard_padded_queries: Caller guarantees masked query rows are unused. - Required with key_padding_mask; padded output rows are zero-filled. + 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: + if key_padding_mask is not None and discard_padded_queries: if ( kv is not None - or not discard_padded_queries or attention_mask != PredefinedAttentionMask.FULL or seq_len_kv not in (None, seq_len) or self.quant_attention_config is not None diff --git a/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/transformer_hunyuan_video1_5.py b/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/transformer_hunyuan_video1_5.py index 87500536631e..0e3ea948fe3b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/transformer_hunyuan_video1_5.py +++ b/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/transformer_hunyuan_video1_5.py @@ -298,21 +298,6 @@ def __init__( bias: bool = True, model_config: Optional[DiffusionModelConfig] = None, ) -> None: - model_config = model_config or DiffusionModelConfig() - if model_config.attention.backend == "TRTLLM": - # The text refiner needs general key padding, not the valid-prefix - # and discarded-query contract of TRTLLM padded self-attention. - model_config = model_config.model_copy( - update={ - "attention": model_config.attention.model_copy( - update={ - "backend": "VANILLA", - "quant_attention_config": None, - "sparse_attention_config": None, - } - ) - } - ) super().__init__( hidden_size=dim, num_attention_heads=num_attention_heads, @@ -320,7 +305,7 @@ def __init__( qkv_mode=QKVMode.FUSE_QKV, qk_norm=False, bias=bias, - config=model_config, + config=model_config or DiffusionModelConfig(), ) def forward( @@ -661,11 +646,7 @@ def forward( attention_mask, (seq_len - attention_mask.shape[1], 0), value=True ).bool() - # Context tokens are packed valid-first; padded context rows remain - # masked in every block and do not contribute to the video output. - hidden_states = self._attn_impl( - query, key, value, key_padding_mask=key_padding_mask, discard_padded_queries=True - ) + hidden_states = self._attn_impl(query, key, value, key_padding_mask=key_padding_mask) # hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) 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 4748f848f408..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,6 +41,8 @@ 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] @@ -69,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/test_hunyuan_video1_5_transformer.py b/tests/unittest/_torch/visual_gen/test_hunyuan_video1_5_transformer.py index de633f1979eb..a3d3ddf3193c 100644 --- a/tests/unittest/_torch/visual_gen/test_hunyuan_video1_5_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_hunyuan_video1_5_transformer.py @@ -281,42 +281,3 @@ def test_hunyuan_video1_5_allclose_to_hf(self): if __name__ == "__main__": pytest.main([__file__, "-v"]) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -@pytest.mark.parametrize("valid_tokens", [[True, True, False, False], [False, True, False, True]]) -def test_trtllm_refiner_respects_key_padding(valid_tokens: list[bool]) -> None: - from tensorrt_llm._torch.visual_gen.attention.backends.vanilla import VanillaAttention - from tensorrt_llm._torch.visual_gen.models.hunyuan_video1_5.transformer_hunyuan_video1_5 import ( - HunyuanVideo15RefinerAttention, - ) - - config = _create_model_config(CONFIG, backend="TRTLLM") - refiner = HunyuanVideo15RefinerAttention(256, 2, 128, model_config=config) - refiner = refiner.to(device="cuda", dtype=torch.bfloat16).eval() - assert isinstance(refiner.attn, VanillaAttention) - assert config.attention.backend == "TRTLLM" - - torch.manual_seed(42) - with torch.no_grad(): - for parameter in refiner.parameters(): - parameter.normal_(std=0.02) - hidden = torch.randn(1, 4, 256, device="cuda", dtype=torch.bfloat16) - mask = torch.tensor([valid_tokens], device="cuda") - q, k, v = refiner.get_qkv(hidden) - expected = ( - F.scaled_dot_product_attention( - q.view(1, 4, 2, 128).transpose(1, 2), - k.view(1, 4, 2, 128).transpose(1, 2), - v.view(1, 4, 2, 128).transpose(1, 2), - attn_mask=mask[:, None, None, :], - ) - .transpose(1, 2) - .reshape(1, 4, 256) - ) - expected = refiner.to_out[0](expected) - actual = refiner(hidden, mask) - torch.testing.assert_close(actual, expected) - hidden[~mask] += 10 - perturbed = refiner(hidden, mask) - torch.testing.assert_close(perturbed[mask], actual[mask]) 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 05e38feda0fb..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", ["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 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 2a518bc68119..a8b2f71f2daf 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -381,7 +381,8 @@ def forward() -> torch.Tensor: @_REQUIRES_SM100 @pytest.mark.parametrize("fused", [False, True]) -def test_trtllm_padded_self_attention_matches_valid_rows(fused: bool) -> None: +@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, @@ -390,8 +391,11 @@ def test_trtllm_padded_self_attention_matches_valid_rows(fused: bool) -> None: attention_metadata_state=create_attention_metadata_state(), ) torch.manual_seed(42) - q, k, v = [torch.randn(2, 8, 2, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3)] - mask = torch.arange(8, device="cuda")[None] < torch.tensor([5, 7], device="cuda")[:, None] + 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] @@ -400,14 +404,60 @@ def test_trtllm_padded_self_attention_matches_valid_rows(fused: bool) -> None: .flatten(2) ) query = torch.cat((q, k, v), dim=2) if fused else q - kwargs = dict(batch_size=2, seq_len=8, key_padding_mask=mask, discard_padded_queries=True) - actual = attention.forward(query, None if fused else k, None if fused else v, **kwargs) + 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, batch_size=2, seq_len=8, key_padding_mask=mask) + 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) + 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) From fd2f409ba69f53b0862fed5509d8bab1ec494210 Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Fri, 30 Oct 2026 19:23:41 +0000 Subject: [PATCH 4/4] Fix metadata state initialization in padding opt-in test Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- .../_torch/visual_gen/test_trtllm_attention_metadata.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) 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 a8b2f71f2daf..8d9044ac6e0f 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -427,7 +427,12 @@ def test_trtllm_padded_self_attention_matches_valid_rows(fused: bool, lengths: l @_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 = 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"))