revert: remove FA4-only vision metadata changes

Signed-off-by: Ajay Balasa <abalasa@nvidia.com>

Author: balasaajay

src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_model.py

@@ -267,7 +267,7 @@ def build_packed_seq_params(
         cu_seqlens = seqlens.cumsum(dim=0)
         cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).int()
 
-        max_seqlen_q = int(seqlens.max().item())
+        max_seqlen_q = seqlens.max()
         return PackedSeqParams(
             cu_seqlens_q=cu_seqlens,
             cu_seqlens_kv=cu_seqlens,

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/vision_model.py

@@ -383,7 +383,7 @@ def build_packed_seq_params(
         cu_seqlens = seqlens.cumsum(dim=0)
         cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).int()
 
-        max_seqlen_q = int(seqlens.max().item())
+        max_seqlen_q = seqlens.max()
         return PackedSeqParams(
             cu_seqlens_q=cu_seqlens,
             cu_seqlens_kv=cu_seqlens,

src/megatron/bridge/training/nemotron_omni_step.py

@@ -158,7 +158,7 @@ def _build_vision_packed_seq_params(
         cu.append(cu[-1] + sl)
     device = imgs_sizes.device if torch.is_tensor(imgs_sizes) else torch.device("cpu")
     cu_tensor = torch.tensor(cu, dtype=torch.int32, device=device)
-    max_len = max(seq_lens, default=0)
+    max_len = torch.tensor(max(seq_lens) if seq_lens else 0, dtype=torch.int32, device=device)
     return PackedSeqParams(
         qkv_format="thd",
         cu_seqlens_q=cu_tensor,

tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py

@@ -34,30 +34,6 @@
     MultiTypeMoeSubmodules,
 )
 from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.model import ErnieMultimodalRotaryEmbedding
-from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.vision_model import ErnieVLVisionModel
-
-
-@pytest.mark.unit
-@pytest.mark.parametrize(
-    "grid,expected_boundaries,expected_max",
-    [
-        ([[1, 2, 3]], [0, 6], 6),
-        ([[2, 2, 3], [1, 2, 4]], [0, 6, 12, 20], 8),
-    ],
-)
-def test_vision_packed_metadata_uses_integer_maxima(grid, expected_boundaries, expected_max):
-    grid = torch.tensor(grid, dtype=torch.int64)
-
-    packed = ErnieVLVisionModel.build_packed_seq_params(None, grid)
-
-    assert packed.qkv_format == "thd"
-    assert type(packed.max_seqlen_q) is int
-    assert type(packed.max_seqlen_kv) is int
-    assert packed.max_seqlen_q == packed.max_seqlen_kv == expected_max
-    for boundaries in (packed.cu_seqlens_q, packed.cu_seqlens_kv):
-        assert boundaries.tolist() == expected_boundaries
-        assert boundaries.dtype == torch.int32
-        assert boundaries.device == grid.device
 
 
 def _make_vision_config():

tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_vision_model.py

@@ -95,36 +95,6 @@ def test_raises_when_seq_exceeds_max(self):
 class TestVisionForwardPackedAttentionSetup:
     """Tests for ``_vision_forward_packed_attention_setup``."""
 
-    @pytest.mark.unit
-    @pytest.mark.parametrize(
-        "grid,expected_boundaries,expected_max",
-        [
-            ([[1, 2, 3]], [0, 6], 6),
-            ([[2, 2, 3], [1, 2, 4]], [0, 6, 12, 20], 8),
-        ],
-    )
-    def test_non_cuda_graph_path_uses_integer_maxima(self, grid, expected_boundaries, expected_max):
-        grid = torch.tensor(grid, dtype=torch.int64)
-        seq_len = expected_boundaries[-1]
-        packed, mask = _vision_forward_packed_attention_setup(
-            use_cuda_graph_padding=False,
-            hidden_states=torch.zeros(seq_len, 1, 8),
-            original_seq_len=seq_len,
-            seq_len=seq_len,
-            grid_thw=grid,
-            build_packed_seq_params=lambda value: Qwen3VLVisionModel.build_packed_seq_params(None, value),
-        )
-
-        assert mask is None
-        assert packed.qkv_format == "thd"
-        assert type(packed.max_seqlen_q) is int
-        assert type(packed.max_seqlen_kv) is int
-        assert packed.max_seqlen_q == packed.max_seqlen_kv == expected_max
-        for boundaries in (packed.cu_seqlens_q, packed.cu_seqlens_kv):
-            assert boundaries.tolist() == expected_boundaries
-            assert boundaries.dtype == torch.int32
-            assert boundaries.device == grid.device
-
     def test_non_cuda_graph_path_calls_builder_and_no_mask(self):
         grid = torch.tensor([[1, 4, 4]], dtype=torch.long)
         hidden = torch.zeros(5, 1, 8, dtype=torch.float32)

tests/unit_tests/training/test_nemotron_omni_step.py

@@ -22,29 +22,6 @@
 from megatron.bridge.training.utils.packed_seq_utils import get_packed_seq_params
 
 
-@pytest.mark.unit
-@pytest.mark.parametrize(
-    "image_sizes,expected_boundaries,expected_max",
-    [
-        ([[512, 512]], [0, 1024], 1024),
-        ([[32, 48], [64, 32]], [0, 6, 14], 8),
-    ],
-)
-def test_vision_packed_metadata_uses_integer_maxima(image_sizes, expected_boundaries, expected_max):
-    image_sizes = torch.tensor(image_sizes, dtype=torch.int64)
-
-    metadata = nemotron_omni_step._build_vision_packed_seq_params(image_sizes)
-
-    assert metadata.qkv_format == "thd"
-    assert type(metadata.max_seqlen_q) is int
-    assert type(metadata.max_seqlen_kv) is int
-    assert metadata.max_seqlen_q == metadata.max_seqlen_kv == expected_max
-    for boundaries in (metadata.cu_seqlens_q, metadata.cu_seqlens_kv):
-        assert boundaries.tolist() == expected_boundaries
-        assert boundaries.dtype == torch.int32
-        assert boundaries.device == image_sizes.device
-
-
 def test_batch_moves_only_one_compatible_token_alias_to_cuda(monkeypatch):
     cuda_inputs = []
 
@@ -145,27 +122,6 @@ def _pipeline_cfg(*, packed=True, defer_packing=False, temporal_patch_dim=1):
     )
 
 
-@pytest.mark.unit
-@pytest.mark.parametrize("packed", [False, True])
-def test_first_pipeline_batch_keeps_vision_maxima_as_integers(monkeypatch, packed):
-    monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self)
-    monkeypatch.setattr(nemotron_omni_step, "is_pp_first_stage", lambda group: True)
-    monkeypatch.setattr(nemotron_omni_step, "is_pp_last_stage", lambda group: False)
-    batch = _packed_pipeline_batch()
-    if not packed:
-        for key in (*nemotron_omni_step._PACKED_SEQ_PARAM_KEYS, "padding_mask"):
-            batch.pop(key, None)
-
-    result = get_batch(iter([batch]), _pipeline_cfg(packed=packed), pg_collection=SimpleNamespace(pp=object()))
-
-    vision_metadata = result[12]
-    assert type(vision_metadata.max_seqlen_q) is int
-    assert type(vision_metadata.max_seqlen_kv) is int
-    assert vision_metadata.max_seqlen_q == vision_metadata.max_seqlen_kv == 4
-    assert vision_metadata.cu_seqlens_q.tolist() == [0, 4, 8]
-    assert vision_metadata.cu_seqlens_kv.tolist() == [0, 4, 8]
-
-
 def test_forward_rejects_deferred_multimodal_packing():
     state = SimpleNamespace(
         timers=None,