diff --git a/tensorrt_llm/_torch/models/modeling_gemma3.py b/tensorrt_llm/_torch/models/modeling_gemma3.py index 5f323ca2ca03..075f40a6b1ee 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3.py +++ b/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( diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index d5fa9c9c5da2..94999d583813 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/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, diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 536594b3ebdf..afdccd638c3f 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/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()