Fix the test of multimodal manager (vllm-project#986)

Signed-off-by: Kewei Wang <[email protected]>

Author: KWang1998

tests/runner/test_multimodal_manager.py

@@ -304,11 +304,11 @@ def test_gather_mm_embeddings_chunked_prefill(self):
         mock_scheduler_output_1.num_scheduled_tokens = {req_id: 20}
 
         gathered_embeds_1 = self.runner.mm_manager.gather_mm_embeddings(
-            mock_scheduler_output_1)
+            mock_scheduler_output_1, target_pad_len=10)
 
-        assert len(gathered_embeds_1) == 1
         expected_embeds_1 = encoder_embedding[0:10]
-        np.testing.assert_array_equal(np.asarray(gathered_embeds_1[0]),
+        assert gathered_embeds_1.shape == expected_embeds_1.shape
+        np.testing.assert_array_equal(np.asarray(gathered_embeds_1),
                                       np.asarray(expected_embeds_1))
 
         # ----- Step 2: Middle chunk of prefill -----
@@ -317,11 +317,11 @@ def test_gather_mm_embeddings_chunked_prefill(self):
         mock_scheduler_output_2.num_scheduled_tokens = {req_id: 30}
 
         gathered_embeds_2 = self.runner.mm_manager.gather_mm_embeddings(
-            mock_scheduler_output_2)
+            mock_scheduler_output_2, target_pad_len=30)
 
-        assert len(gathered_embeds_2) == 1
         expected_embeds_2 = encoder_embedding[10:40]
-        np.testing.assert_array_equal(np.asarray(gathered_embeds_2[0]),
+        assert gathered_embeds_2.shape == expected_embeds_2.shape
+        np.testing.assert_array_equal(np.asarray(gathered_embeds_2),
                                       np.asarray(expected_embeds_2))
 
         # ----- Step 3: Last chunk of prefill -----
@@ -330,11 +330,11 @@ def test_gather_mm_embeddings_chunked_prefill(self):
         mock_scheduler_output_3.num_scheduled_tokens = {req_id: 30}
 
         gathered_embeds_3 = self.runner.mm_manager.gather_mm_embeddings(
-            mock_scheduler_output_3)
+            mock_scheduler_output_3, target_pad_len=16)
 
-        assert len(gathered_embeds_3) == 1
         expected_embeds_3 = encoder_embedding[40:56]
-        np.testing.assert_array_equal(np.asarray(gathered_embeds_3[0]),
+        assert gathered_embeds_3.shape == expected_embeds_3.shape
+        np.testing.assert_array_equal(np.asarray(gathered_embeds_3),
                                       np.asarray(expected_embeds_3))
 
     def test_calc_mrope_positions(self):