fix(nemotron-omni): accept dense decoder masks

src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py

@@ -292,7 +292,19 @@ def _merge_projected_media(
 
         media_mask = input_ids == media_token_id
         if attention_mask is not None:
-            media_mask = media_mask & attention_mask.bool()
+            if attention_mask.ndim == input_ids.ndim:
+                if attention_mask.shape != input_ids.shape:
+                    raise ValueError(
+                        "A token-validity attention mask must match input_ids; "
+                        f"got {tuple(attention_mask.shape)} and {tuple(input_ids.shape)}."
+                    )
+                media_mask = media_mask & attention_mask.bool()
+            elif attention_mask.ndim != 4:
+                raise ValueError(
+                    "Nemotron Omni expects either a token-validity mask with the "
+                    "same rank as input_ids or a four-dimensional decoder mask; "
+                    f"got attention_mask.ndim={attention_mask.ndim}."
+                )
 
         expected_features = int(media_mask.sum().item())
         actual_features = media_embeddings.shape[0]

tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py

@@ -327,6 +327,24 @@ def test_padded_placeholder_is_not_treated_as_media():
     assert torch.equal(output[3, 0], torch.zeros(3))
 
 
+def test_media_merge_accepts_dense_decoder_attention_mask():
+    language_embeddings = torch.zeros(4, 2, 3)
+    input_ids = torch.tensor([[7, 18, 9, 0], [8, 18, 10, 0]])
+    attention_mask = torch.tril(torch.ones(2, 1, 4, 4, dtype=torch.bool))
+    media_embeddings = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+
+    output = NemotronOmniModel._merge_projected_media(
+        language_embeddings,
+        input_ids,
+        media_embeddings,
+        media_token_id=18,
+        attention_mask=attention_mask,
+    )
+
+    assert torch.equal(output[1, 0], media_embeddings[0])
+    assert torch.equal(output[1, 1], media_embeddings[1])
+
+
 def test_media_merge_supports_backward_for_batch_size_one():
     language_embeddings = torch.randn(4, 1, 3, requires_grad=True)
     media_embeddings = torch.randn(2, 3, requires_grad=True)