[TRTLLM-17109][perf] Reduce Gemma multimodal attention mask construction peak memory

tensorrt_llm/_torch/models/modeling_gemma3.py

@@ -17,6 +17,8 @@
 from typing import Any, Dict, Literal, Optional, Tuple
 
 import torch
+import triton
+import triton.language as tl
 from torch import nn
 from transformers import Gemma3TextConfig
 
@@ -44,6 +46,87 @@
                              register_auto_model)
 
 
+@triton.jit
+def _write_mask(
+    output,
+    blob_ids,
+    query_length,
+    prefix_length,
+    WINDOW: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    key_length = prefix_length + query_length
+    valid = col < key_length
+    offsets = row.to(tl.int64) * key_length + col
+    allowed = col <= prefix_length + row
+    if WINDOW is not None:
+        allowed &= col > prefix_length + row - WINDOW
+    query_blob = tl.load(blob_ids + row)
+    key_blob = tl.load(blob_ids + col - prefix_length,
+                       valid & (col >= prefix_length),
+                       other=0)
+    allowed |= (query_blob != 0) & (query_blob == key_blob)
+    tl.store(output + offsets, allowed, valid)
+
+
+def build_gemma_attention_mask(
+    token_types: torch.Tensor,
+    context_lengths: list[int],
+    prefix_lengths: list[int],
+    sliding_window: int | None = None,
+) -> torch.Tensor:
+    """Build flattened bool masks for packed Gemma context requests.
+
+    CUDA graph capture requires prior warmup. Request lengths, prefix lengths,
+    and the sliding window are host-side launch parameters and must remain
+    fixed during replay; token-type values may change in place. This function
+    does not modify token_types.
+
+    Args:
+        token_types: One-dimensional CUDA bool image flags (Gemma3) or integer
+            modality IDs (Gemma4: 0=text, 1=image, 2=video, 3=audio). Contiguous
+            runs of the same vision type attend bidirectionally. A strided
+            input is supported; trailing generation tokens are ignored.
+        context_lengths: Number of new tokens for each context request.
+        prefix_lengths: Cached token counts in the same request order. Chunk
+            boundaries must not split a vision blob, as in the model caller.
+        sliding_window: Exclusive lower attention-window bound, or None.
+
+    Returns:
+        Contiguous mask with sum(q * (p + q)) elements. Prefix columns apply
+        causal/window rules; vision blobs override those rules only among
+        new tokens. Intermediate storage is linear in the context length.
+    """
+    mask_sizes = [
+        q * (p + q)
+        for q, p in zip(context_lengths, prefix_lengths, strict=True)
+    ]
+    output = torch.empty(sum(mask_sizes),
+                         dtype=torch.bool,
+                         device=token_types.device)
+    token_offset = 0
+    mask_offset = 0
+    for q, p, size in zip(context_lengths,
+                          prefix_lengths,
+                          mask_sizes,
+                          strict=True):
+        types = token_types[token_offset:token_offset + q]
+        if size:
+            vision = (types == 1) | (types == 2)
+            previous = torch.cat((types.new_zeros(1), types[:-1]))
+            starts = (types != previous) & vision
+            blobs = torch.cumsum(starts, dim=0, dtype=torch.int32)
+            blobs.masked_fill_(~vision, 0)
+            _write_mask[(q, triton.cdiv(
+                p + q, 1024))](output[mask_offset:mask_offset + size], blobs, q,
+                               p, sliding_window, 1024)
+        token_offset += q
+        mask_offset += size
+    return output
+
+
 class Gemma3TextScaledWordEmbedding(Embedding):
 
     def __init__(
@@ -321,82 +404,18 @@ def get_preferred_kv_cache_manager_version(cls,
         """
         return "V2"
 
-    def _get_token_type_mask(self, image_token_mask: torch.BoolTensor):
-        device = image_token_mask.device
-        sequence_length = len(image_token_mask)
-
-        # Create a list of token type ids. 0 for text tokens, 1 for all image tokens (regardless of which image they belong to).
-        token_type_ids = torch.zeros(sequence_length,
-                                     dtype=torch.int32,
-                                     device=device)
-        token_type_ids[image_token_mask] = 1
-
-        # There could be image tokens from multiple images where those corresponding to the
-        # same image are contiguous. We assign a unique id to each contiguous blob of image tokens now.
-
-        # Pad with zero at the start to detect changes.
-        padded = torch.cat((torch.tensor([0], device=device), token_type_ids))
-
-        # Identify where blobs start (0->1 transitions)
-        starts = (padded[1:] > padded[:-1]).int()
-
-        # Cumulative sum of starts gives a unique id for each blob. Note that
-        # this assigns a unique id to the zeros separating the blobs.
-        blob_ids = torch.cumsum(starts, dim=0)
-
-        # Mask out zeros (positions where token_type_ids == 0).
-        token_type_ids *= blob_ids
-
-        # Create a mask where each blob is a unique id.
-        token_type_mask = token_type_ids.unsqueeze(
-            0) == token_type_ids.unsqueeze(1)
-
-        # If text token, do not change anything.
-        token_type_mask = torch.where(token_type_ids == 0, False,
-                                      token_type_mask)
-
-        return token_type_mask
-
     def get_context_mask(
         self,
         image_token_mask: torch.BoolTensor,
         effective_sliding_window: Optional[int] = None,
-    ):
-        """
-        Returns an attention mask such that text tokens attend to each other in causal fashion while image
-        tokens attend in causal fashion as well as to all other image tokens in a bidirectional manner.
-        Args:
-            image_token_mask: A boolean tensor of shape (sequence_length,) where True indicates an image token.
-            effective_sliding_window: The effective sliding window size for the attention mask. Default is None, which means no sliding window.
-            For Gemma3, this is the sliding window size from config (e.g. 512 for 1B model).
-        Returns:
-            A boolean attention mask of shape (sequence_length, sequence_length).
-        """
-        device = image_token_mask.device
-        sequence_length = len(image_token_mask)
-        if effective_sliding_window is None or effective_sliding_window >= sequence_length:
-            causal_mask = torch.arange(
-                sequence_length, device=device).unsqueeze(0) <= torch.arange(
-                    sequence_length, device=device).unsqueeze(1)
-        else:
-            attention_mask_1 = (torch.arange(sequence_length,
-                                             device=device).unsqueeze(0)
-                                <= torch.arange(sequence_length,
-                                                device=device).unsqueeze(1))
-            attention_mask_2 = (
-                torch.arange(sequence_length, device=device).unsqueeze(0)
-                > torch.arange(sequence_length, device=device).unsqueeze(1) -
-                effective_sliding_window)
-            causal_mask = attention_mask_1 & attention_mask_2
-
-        # Apply a bidirectional mask for image tokens.
-        token_type_mask = self._get_token_type_mask(image_token_mask)
-        causal_mask = causal_mask.masked_fill(token_type_mask, True)
-        return causal_mask
-
-    # ASSUMPTIONS:
-    # 1) Chunked prefill is disabled to avoid chunking image tokens as they need bidirectional attention.
-    # 2) KV cache reuse is disabled to avoid partially matched image tokens (entire image must be reused to get things correct).
+    ) -> torch.Tensor:
+        """Build causal/window attention with bidirectional image blobs."""
+        length = len(image_token_mask)
+        return build_gemma_attention_mask(image_token_mask, [length], [0],
+                                          effective_sliding_window).view(
+                                              length, length)
+
+    # Image blobs cannot be split by chunked prefill or partial prefix reuse.
     def get_attention_mask(
             self,
             image_token_mask: torch.BoolTensor,
@@ -436,17 +455,9 @@ def get_attention_mask(
             "and kv cache reuse must be disabled.")
 
         # Create masks for context requests.
-        context_mask_list = []
-        token_offset = 0
-        for context_len in context_lens:
-            mask_i = self.get_context_mask(
-                image_token_mask=image_token_mask[token_offset:token_offset +
-                                                  context_len],
-                effective_sliding_window=effective_sliding_window,
-            )
-            context_mask_list.append(mask_i.flatten())
-            token_offset += context_len
-        return torch.cat(context_mask_list, dim=0).contiguous()
+        return build_gemma_attention_mask(image_token_mask, context_lens,
+                                          cached_token_lens,
+                                          effective_sliding_window)
 
     @inference_mode_unless_compiling
     def forward(

tensorrt_llm/_torch/models/modeling_gemma4.py

@@ -65,6 +65,7 @@
 from ..modules.rms_norm import RMSNorm
 from ..speculative.interface import SpecMetadata
 from ..utils import ActivationType, Fp4QuantizedTensor, is_torch_compiling
+from .modeling_gemma3 import build_gemma_attention_mask
 from .modeling_speculative import (
     _SPECULATIVE_POSITION_HEADROOM,
     SpecDecOneEngineForCausalLM,
@@ -1330,86 +1331,17 @@ def get_preferred_transceiver_runtime(
         """Prefer the Python transceiver so disaggregated serving over NIXL keeps V2."""
         return "PYTHON"
 
-    def _get_token_type_mask(self, mm_token_type_ids: torch.Tensor):
-        """Build bidirectional attention mask from mm_token_type_ids.
-
-        mm_token_type_ids: 0=text, 1=image, 2=video, 3=audio. Only VISION
-        tokens (image/video) attend bidirectionally within their contiguous
-        blob; text and audio stay causal. Matches HF Gemma4, where
-        ``is_vision = (mm_token_type_ids == 1) | (mm_token_type_ids == 2)`` and
-        audio is left causal.
-        """
-        device = mm_token_type_ids.device
-        token_type_ids = mm_token_type_ids.clone()
-        # Only vision tokens (image=1, video=2) get the bidirectional mask;
-        # audio (3) stays causal, matching HF Gemma4 (audio is not in is_vision).
-        is_mm = (token_type_ids == 1) | (token_type_ids == 2)
-
-        # Detect blob boundaries: positions where type changes or goes 0->nonzero
-        padded = torch.cat(
-            (
-                torch.tensor([0], device=device, dtype=mm_token_type_ids.dtype),
-                mm_token_type_ids,
-            )
-        )
-        # A new blob starts whenever the type changes and the current is nonzero
-        changes = (padded[1:] != padded[:-1]) & is_mm
-        blob_ids = torch.cumsum(changes.int(), dim=0)
-        # Mask out text tokens
-        blob_ids = blob_ids * is_mm.int()
-
-        # Tokens with the same non-zero blob_id attend bidirectionally
-        token_type_mask = blob_ids.unsqueeze(0) == blob_ids.unsqueeze(1)
-        token_type_mask.logical_and_(is_mm.unsqueeze(0))
-        token_type_mask.logical_and_(is_mm.unsqueeze(1))
-
-        return token_type_mask
-
     def get_context_mask(
         self,
         mm_token_type_ids: torch.Tensor,
         effective_sliding_window: Optional[int] = None,
         prefix_len: int = 0,
-    ):
-        """Build context mask with causal + bidirectional for MM tokens.
-
-        Returns a [extend_len, prefix_len + extend_len] mask where:
-        - The first `prefix_len` columns (cached/paged history) apply the
-          sliding window using absolute token positions. Bidirectional MM
-          across the prefix/extend boundary is NOT supported here; callers
-          must ensure chunk boundaries do not split a multimodal block.
-        - The last `extend_len` columns follow the original causal +
-          (optional) sliding window + MM-bidirectional logic.
-        """
-        device = mm_token_type_ids.device
-        extend_len = len(mm_token_type_ids)
-
-        pos = torch.arange(extend_len, device=device)
-        causal_mask = pos.unsqueeze(0) <= pos.unsqueeze(1)
-
-        if effective_sliding_window is not None and effective_sliding_window < extend_len:
-            causal_mask.logical_and_(pos.unsqueeze(0) > pos.unsqueeze(1) - effective_sliding_window)
-
-        token_type_mask = self._get_token_type_mask(mm_token_type_ids)
-        causal_mask = causal_mask.masked_fill(token_type_mask, True)
-
-        if prefix_len > 0:
-            if (
-                effective_sliding_window is not None
-                and effective_sliding_window < prefix_len + extend_len
-            ):
-                query_pos = prefix_len + pos
-                prefix_pos = torch.arange(prefix_len, device=device)
-                prefix_block = (
-                    prefix_pos.unsqueeze(0) > query_pos.unsqueeze(1) - effective_sliding_window
-                )
-            else:
-                prefix_block = torch.ones(
-                    extend_len, prefix_len, dtype=causal_mask.dtype, device=device
-                )
-            causal_mask = torch.cat([prefix_block, causal_mask], dim=1)
-
-        return causal_mask
+    ) -> torch.Tensor:
+        """Build a context mask; chunk boundaries must not split vision blobs."""
+        length = len(mm_token_type_ids)
+        return build_gemma_attention_mask(
+            mm_token_type_ids, [length], [prefix_len], effective_sliding_window
+        ).view(length, prefix_len + length)
 
     def get_attention_mask(
         self,
@@ -1432,18 +1364,9 @@ def get_attention_mask(
         else:
             cached_token_lens = [0] * num_contexts
 
-        context_mask_list = []
-        token_offset = 0
-        for context_len, prefix_len in zip(context_lens, cached_token_lens, strict=True):
-            context_end = token_offset + context_len
-            mask_i = self.get_context_mask(
-                mm_token_type_ids=mm_token_type_ids[token_offset:context_end],
-                effective_sliding_window=effective_sliding_window,
-                prefix_len=prefix_len,
-            )
-            context_mask_list.append(mask_i.flatten())
-            token_offset = context_end
-        return torch.cat(context_mask_list, dim=0).contiguous()
+        return build_gemma_attention_mask(
+            mm_token_type_ids, context_lens, cached_token_lens, effective_sliding_window
+        )
 
     def _forward_speculative(
         self,

tests/unittest/_torch/modeling/test_modeling_gemma4.py

@@ -35,6 +35,7 @@
 from tensorrt_llm._torch.metadata import KVCacheParams
 from tensorrt_llm._torch.model_config import ModelConfig
 from tensorrt_llm._torch.models.checkpoints.hf.gemma4_weight_mapper import Gemma4HfWeightMapper
+from tensorrt_llm._torch.models.modeling_gemma3 import build_gemma_attention_mask
 from tensorrt_llm._torch.models.modeling_gemma4 import (
     Gemma4AssistantForCausalLM,
     Gemma4AssistantMaskedEmbedder,
@@ -2404,6 +2405,15 @@ def test_chunked_context_mask_applies_prefix_window(self) -> None:
         self.assertFalse(chunk_mask[-1, 0].item())
         self.assertTrue(chunk_mask[-1, 32].item())
 
+        # Vision runs override the two-token window; audio and text stay causal.
+        mixed_types = torch.tensor([1, 1, 1, 2, 2, 2, 3, 3, 3, 0], device="cuda")
+        mixed_mask = model.get_context_mask(mixed_types, effective_sliding_window=2)
+        expected = torch.eye(10, dtype=torch.bool)
+        expected[torch.arange(1, 10), torch.arange(9)] = True
+        expected[:3, :3] = True
+        expected[3:6, 3:6] = True
+        torch.testing.assert_close(mixed_mask.cpu(), expected, rtol=0, atol=0)
+
     @torch.no_grad()
     def test_bidirectional_mask_only_applies_to_sliding_layers(self):
         """Full-attention layers retain the standard causal mask."""
@@ -4373,5 +4383,23 @@ def test_vision_attn_metadata_max_num_requests_is_not_one(self):
         )
 
 
+def test_gemma_attention_mask_peak() -> None:
+    # A second quadratic tensor exceeds the allowance by tens of MiB.
+    length = 4096
+    types = torch.zeros(2 * length, dtype=torch.int64, device="cuda")
+    types[64:128] = 1
+    warm = build_gemma_attention_mask(types, [length, length], [0, 1024], 512)
+    del warm
+    torch.cuda.synchronize()
+    before = torch.cuda.memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
+    result = build_gemma_attention_mask(types, [length, length], [0, 1024], 512)
+    torch.cuda.synchronize()
+    peak = torch.cuda.max_memory_allocated() - before
+    assert result.dtype == torch.bool
+    assert result.numel() == length * length + length * (1024 + length)
+    assert peak <= result.numel() + 2 * 1024**2
+
+
 if __name__ == "__main__":
     unittest.main()