diff --git a/tensorrt_llm/_torch/models/modeling_gemma3.py b/tensorrt_llm/_torch/models/modeling_gemma3.py
index f13a626bb359..460af83ddbff 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma3.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma3.py
@@ -48,6 +48,7 @@ def __init__(
         )
         self.embed_scale = torch.sqrt(torch.tensor(hidden_size)).to(self.dtype)
 
+    @torch.inference_mode()
     def forward(self, input_ids):
         return super().forward(input_ids) * self.embed_scale
 
@@ -97,6 +98,7 @@ def __init__(
         self.aux_stream = torch.cuda.Stream()
         self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
 
+    @torch.inference_mode()
     def forward(
         self,
         position_ids: Optional[torch.IntTensor],
@@ -176,6 +178,7 @@ def __init__(self, config: Gemma3TextConfig):
                                 dtype=self.dtype)
         self.act_fn = ACT2FN[config.hidden_activation]
 
+    @torch.inference_mode()
     def forward(self, x):
         down_proj = self.down_proj(
             self.act_fn(self.gate_proj(x)) * self.up_proj(x))
@@ -215,6 +218,7 @@ def __init__(
             eps=config.rms_norm_eps,
             dtype=config.torch_dtype)
 
+    @torch.inference_mode()
     def forward(
         self,
         position_ids: torch.IntTensor,
@@ -272,6 +276,7 @@ def __init__(self, model_config: ModelConfig[Gemma3TextConfig]):
                             eps=config.pretrained_config.rms_norm_eps,
                             dtype=config.pretrained_config.torch_dtype)
 
+    @torch.inference_mode()
     def forward(
         self,
         attn_metadata: AttentionMetadata,
@@ -318,6 +323,42 @@ def __init__(
                          hidden_size=model_config.pretrained_config.hidden_size,
                          vocab_size=model_config.pretrained_config.vocab_size)
 
+    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,
@@ -351,15 +392,7 @@ def get_context_mask(
             causal_mask = attention_mask_1 & attention_mask_2
 
         # Apply a bidirectional mask for image tokens.
-        token_type_ids = torch.zeros(sequence_length,
-                                     dtype=torch.int32,
-                                     device=device)
-        # 1 for image tokens, 0 for text tokens.
-        token_type_ids[image_token_mask] = 1
-        token_type_mask = token_type_ids.unsqueeze(
-            0) == token_type_ids.unsqueeze(1)
-        # If text token, do not change anything.
-        token_type_mask[token_type_ids == 0] = False
+        token_type_mask = self._get_token_type_mask(image_token_mask)
         causal_mask = causal_mask.masked_fill(token_type_mask, True)
         return causal_mask
 
@@ -411,6 +444,7 @@ def get_flashinfer_attention_mask(
             context_mask_list.append(mask_i.flatten())
         return torch.cat(context_mask_list, dim=0).contiguous()
 
+    @torch.inference_mode()
     def forward(
         self,
         attn_metadata: AttentionMetadata,
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma3.py b/tests/unittest/_torch/modeling/test_modeling_gemma3.py
index d6228f60f321..d7e27de24b5e 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma3.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma3.py
@@ -385,6 +385,8 @@ def test_gemma3_flashinfer_mask(self):
                                  dtype=torch.int,
                                  device=device)
 
+        # This initial setup is to populate KV cache so that attn_metadata has the info
+        # needed for generating mask.
         context_sequence_lengths = [3, 2, 1]
         sequence_lengths = context_sequence_lengths + [1, 1]
         past_seen_tokens = [0, 0, 0, 2, 1]
@@ -423,19 +425,17 @@ def test_gemma3_flashinfer_mask(self):
         )
         attn_metadata.prepare()
 
-        # Test for case where all tokens are text tokens. `test_gemma3_global_mask` and `test_gemma3_local_mask`
-        # test for individual requests where there's a mix of text and image tokens.
+        # First sample has 2 image tokens, second sample has 2 image tokens, third sample has none.
         image_token_mask = torch.tensor(
-            [False, False, False, False, False, False, False, False],
-            device=device)
+            [True, True, False, True, True, True, False, False], device=device)
         causal_mask = gemma3.get_flashinfer_attention_mask(
             image_token_mask=image_token_mask, attn_metadata=attn_metadata)
         # Causal mask for context request 1.
         ctx_request_1_mask = torch.tensor(
-            [[True, False, False], [True, True, False], [True, True, True]],
+            [[True, True, False], [True, True, False], [True, True, True]],
             device=device)
         # Causal mask for context request 2.
-        ctx_request_2_mask = torch.tensor([[True, False], [True, True]],
+        ctx_request_2_mask = torch.tensor([[True, True], [True, True]],
                                           device=device)
         # Causal mask for context request 3.
         ctx_request_3_mask = torch.tensor([[True]], device=device)
@@ -448,7 +448,8 @@ def test_gemma3_flashinfer_mask(self):
                                          dim=0)
         torch.testing.assert_close(causal_mask, expected_causal_mask)
 
-    def test_gemma3_global_context_mask(self) -> None:
+    def single_image_mask_test_helper(
+            self, effective_sliding_window: int) -> torch.Tensor:
         config_dict = deepcopy(GEMMA3_1B_CONFIG)
         gemma3_config = Gemma3TextConfig.from_dict(config_dict)
         device = torch.device('cuda')
@@ -458,34 +459,54 @@ def test_gemma3_global_context_mask(self) -> None:
 
         image_token_mask = torch.tensor(
             [False, False, True, True, True, True, False, False], device=device)
-        attention_mask = gemma3.get_context_mask(
-            image_token_mask=image_token_mask, effective_sliding_window=None)
+        return gemma3.get_context_mask(
+            image_token_mask=image_token_mask,
+            effective_sliding_window=effective_sliding_window)
 
+    @parameterized.expand([
+        "global",
+        "sliding_window_larger_than_seq",
+        "sliding_window_equal_to_seq",
+    ], lambda testcase_func, param_num, param:
+                          f"{testcase_func.__name__}[{param.args[0]}]")
+    def test_gemma3_global_context_mask_single_image(self,
+                                                     test_name: str) -> None:
+        device = torch.device('cuda')
+        effecive_sliding_window_map = {
+            # For global mask, don't mention sliding window size.
+            "global": None,
+            # For a sliding window larger than sequence length, the local mask is the same as the global mask.
+            "sliding_window_larger_than_seq": 10,
+            # For a sliding window same as sequence length, the local mask is the same as the global mask.
+            "sliding_window_equal_to_seq": 8,
+        }
+        effective_sliding_window = effecive_sliding_window_map[test_name]
+        attention_mask = self.single_image_mask_test_helper(
+            effective_sliding_window=effective_sliding_window)
         # Text tokens attend to each other in causal fashion. Image tokens attend in causal fashion
         # as well as to all other image tokens.
         expected_attention_mask = torch.tensor(
-            [[True, False, False, False, False, False, False, False],
-             [True, True, False, False, False, False, False, False],
-             [True, True, True, True, True, True, False, False],
-             [True, True, True, True, True, True, False, False],
-             [True, True, True, True, True, True, False, False],
-             [True, True, True, True, True, True, False, False],
-             [True, True, True, True, True, True, True, False],
-             [True, True, True, True, True, True, True, True]],
-            device=device)
-        torch.testing.assert_close(attention_mask, expected_attention_mask)
-
-        # Even for a sliding window larger than sequence length, the attention mask is the same as the global mask.
-        attention_mask = gemma3.get_context_mask(
-            image_token_mask=image_token_mask, effective_sliding_window=10)
+            [[1, 0, 0, 0, 0, 0, 0, 0], [1, 1, 0, 0, 0, 0, 0, 0],
+             [1, 1, 1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 1, 1, 0, 0],
+             [1, 1, 1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 1, 1, 0, 0],
+             [1, 1, 1, 1, 1, 1, 1, 0], [1, 1, 1, 1, 1, 1, 1, 1]],
+            device=device).bool()
         torch.testing.assert_close(attention_mask, expected_attention_mask)
 
-        # Even for sliding window same as sequence length, the attention mask is the same as the global mask.
-        attention_mask = gemma3.get_context_mask(
-            image_token_mask=image_token_mask, effective_sliding_window=8)
+    def test_gemma3_local_context_mask_single_image(self) -> None:
+        device = torch.device('cuda')
+        attention_mask = self.single_image_mask_test_helper(
+            effective_sliding_window=2)
+        expected_attention_mask = torch.tensor(
+            [[1, 0, 0, 0, 0, 0, 0, 0], [1, 1, 0, 0, 0, 0, 0, 0],
+             [0, 1, 1, 1, 1, 1, 0, 0], [0, 0, 1, 1, 1, 1, 0, 0],
+             [0, 0, 1, 1, 1, 1, 0, 0], [0, 0, 1, 1, 1, 1, 0, 0],
+             [0, 0, 0, 0, 0, 1, 1, 0], [0, 0, 0, 0, 0, 0, 1, 1]],
+            device=device).bool()
         torch.testing.assert_close(attention_mask, expected_attention_mask)
 
-    def test_gemma3_local_context_mask(self) -> None:
+    def multi_image_mask_test_helper(
+            self, effective_sliding_window: int) -> torch.Tensor:
         config_dict = deepcopy(GEMMA3_1B_CONFIG)
         gemma3_config = Gemma3TextConfig.from_dict(config_dict)
         device = torch.device('cuda')
@@ -493,18 +514,140 @@ def test_gemma3_local_context_mask(self) -> None:
                                    attn_backend="FLASHINFER")
         gemma3 = Gemma3ForCausalLM(model_config).to(device)
 
+        # 4 images with 4, 3, 2, 2 tokens respectively.
         image_token_mask = torch.tensor(
-            [False, False, True, True, True, True, False, False], device=device)
-        attention_mask = gemma3.get_context_mask(
-            image_token_mask=image_token_mask, effective_sliding_window=2)
-        expected_attention_mask = torch.tensor(
-            [[True, False, False, False, False, False, False, False],
-             [True, True, False, False, False, False, False, False],
-             [False, True, True, True, True, True, False, False],
-             [False, False, True, True, True, True, False, False],
-             [False, False, True, True, True, True, False, False],
-             [False, False, True, True, True, True, False, False],
-             [False, False, False, False, False, True, True, False],
-             [False, False, False, False, False, False, True, True]],
+            [
+                # text blob.
+                False,
+                False,
+                # image1 blob.
+                True,
+                True,
+                True,
+                True,
+                # text blob.
+                False,
+                False,
+                # image2 blob.
+                True,
+                True,
+                True,
+                # text blob.
+                False,
+                False,
+                # image3 blob.
+                True,
+                True,
+                # text blob.
+                False,
+                False,
+                # image4 blob.
+                True,
+                True,
+                # text blob.
+                False,
+            ],
             device=device)
+        return gemma3.get_context_mask(
+            image_token_mask=image_token_mask,
+            effective_sliding_window=effective_sliding_window)
+
+    @parameterized.expand([
+        "global",
+        "sliding_window_larger_than_seq",
+        "sliding_window_equal_to_seq",
+    ], lambda testcase_func, param_num, param:
+                          f"{testcase_func.__name__}[{param.args[0]}]")
+    def test_gemma3_global_context_mask_multi_image(self,
+                                                    test_name: str) -> None:
+        device = torch.device('cuda')
+        effecive_sliding_window_map = {
+            # Don't mention sliding window size for global mask.
+            "global": None,
+            # For a sliding window larger than sequence length, the local mask is the same as the global mask.
+            "sliding_window_larger_than_seq": 25,
+            # For a sliding window same as sequence length, the local mask is the same as the global mask.
+            "sliding_window_equal_to_seq": 20,
+        }
+        effective_sliding_window = effecive_sliding_window_map[test_name]
+
+        # Text tokens attend to each other in causal fashion. Image tokens attend in causal fashion
+        # as well as to all other image tokens.
+        expected_attention_mask = torch.tensor(
+            [
+                # text blob.
+                [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                # image1 blob.
+                [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                # text blob.
+                [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                # image2 blob.
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                # text blob.
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
+                # image3 blob.
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
+                # text blob.
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
+                # image4 blob.
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
+                # text blob.
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+            ],
+            device=device).bool()
+        attention_mask = self.multi_image_mask_test_helper(
+            effective_sliding_window=effective_sliding_window)
+        torch.testing.assert_close(attention_mask, expected_attention_mask)
+
+    def test_gemma3_local_context_mask_multi_image(self) -> None:
+        device = torch.device('cuda')
+        attention_mask = self.multi_image_mask_test_helper(
+            effective_sliding_window=3)
+
+        # Text tokens attend to each other in causal fashion. Image tokens attend in causal fashion
+        # as well as to all other image tokens.
+        expected_attention_mask = torch.tensor(
+            [
+                # text blob.
+                [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                # image1 blob.
+                [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                # text blob.
+                [0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                # image2 blob.
+                [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                # text blob.
+                [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
+                # image3 blob.
+                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0],
+                # text blob.
+                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0],
+                # image4 blob.
+                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0],
+                # text blob.
+                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1],
+            ],
+            device=device).bool()
         torch.testing.assert_close(attention_mask, expected_attention_mask)