[VLM] avoid extra cuda-ipc staging for preprocessed input (#26096)

2026-05-24 19:48:06 +08:00
parent 6447596501
commit b6f71d5850
3 changed files with 127 additions and 79 deletions
@@ -45,6 +45,58 @@ class TestMultimodalInputsFromDict(unittest.TestCase):
        self.assertTrue(torch.equal(mm_inputs.mm_items[0].feature, feature_tensor))
        proxy_feature.reconstruct_on_target_device.assert_called_once_with(0)

+    def test_materialize_precomputed_embedding_proxy_without_feature(self):
+        embedding_tensor = torch.tensor([[1.0, 2.0]], dtype=torch.float32)
+        proxy_embedding = _make_proxy_with_reconstruct_result(embedding_tensor)
+        mm_item = MultimodalDataItem(
+            modality=Modality.IMAGE,
+            offsets=[(0, 1)],
+            precomputed_embeddings=proxy_embedding,
+        )
+
+        with (
+            patch.object(schedule_batch.torch.cuda, "is_available", return_value=True),
+            patch.object(schedule_batch.torch.cuda, "current_device", return_value=0),
+            patch.object(
+                schedule_batch.envs.SGLANG_MM_BUFFER_SIZE_MB, "get", return_value=0
+            ),
+        ):
+            mm_inputs = MultimodalInputs.from_dict({"mm_items": [mm_item]})
+
+        self.assertTrue(
+            torch.equal(
+                mm_inputs.mm_items[0].precomputed_embeddings,
+                embedding_tensor,
+            )
+        )
+        proxy_embedding.reconstruct_on_target_device.assert_called_once_with(0)
+
+    def test_materialize_model_specific_proxy_without_feature(self):
+        grid_tensor = torch.tensor([[1, 2, 3]], dtype=torch.int64)
+        proxy_grid = _make_proxy_with_reconstruct_result(grid_tensor)
+        mm_item = MultimodalDataItem(
+            modality=Modality.IMAGE,
+            offsets=[(0, 1)],
+            model_specific_data={"image_grid_thw": proxy_grid},
+        )
+
+        with (
+            patch.object(schedule_batch.torch.cuda, "is_available", return_value=True),
+            patch.object(schedule_batch.torch.cuda, "current_device", return_value=0),
+            patch.object(
+                schedule_batch.envs.SGLANG_MM_BUFFER_SIZE_MB, "get", return_value=0
+            ),
+        ):
+            mm_inputs = MultimodalInputs.from_dict({"mm_items": [mm_item]})
+
+        self.assertTrue(
+            torch.equal(
+                mm_inputs.mm_items[0].model_specific_data["image_grid_thw"],
+                grid_tensor,
+            )
+        )
+        proxy_grid.reconstruct_on_target_device.assert_called_once_with(0)
+

 if __name__ == "__main__":
    unittest.main(verbosity=2)