Skip to content

enh: Bidirectional mask with multiple images for Gemma3 #5976

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
byshiue merged 1 commit into from
Jul 14, 2025
Merged

enh: Bidirectional mask with multiple images for Gemma3 #5976

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 43 additions & 9 deletions tensorrt_llm/_torch/models/modeling_gemma3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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):
Comment thread
brb-nv marked this conversation as resolved.
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,
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand Down
221 changes: 182 additions & 39 deletions tests/unittest/_torch/modeling/test_modeling_gemma3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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.
Comment thread
brb-nv marked this conversation as resolved.
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)
Expand All @@ -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')
Expand All @@ -458,53 +459,195 @@ 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')
model_config = ModelConfig(pretrained_config=gemma3_config,
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)