diff --git a/tests/multimodal/test_inputs.py b/tests/multimodal/test_inputs.py index d6bdf76a6f7..7752a543f42 100644 --- a/tests/multimodal/test_inputs.py +++ b/tests/multimodal/test_inputs.py @@ -26,11 +26,8 @@ def test_placeholder_range_get_num_embeds(is_embed, expected): "is_embed,expected", [ (None, None), - ( - torch.tensor([False, True, False, True, True]), - torch.tensor([0, 1, 1, 2, 3]), - ), - (torch.tensor([True, True, True]), torch.tensor([1, 2, 3])), + (torch.tensor([False, True, False, True, True]), [0, 1, 1, 2, 3]), + (torch.tensor([True, True, True]), [1, 2, 3]), ], ) def test_placeholder_range_embeds_cumsum(is_embed, expected): @@ -41,6 +38,6 @@ def test_placeholder_range_embeds_cumsum(is_embed, expected): assert pr.embeds_cumsum is None return - assert torch.equal(pr.embeds_cumsum, expected) + assert pr.embeds_cumsum == expected # cached_property should return the same object on repeated access assert pr.embeds_cumsum is pr.embeds_cumsum diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py index 12356b8727c..d98a1624ac3 100644 --- a/vllm/multimodal/inputs.py +++ b/vllm/multimodal/inputs.py @@ -145,14 +145,15 @@ class PlaceholderRange: """ @cached_property - def embeds_cumsum(self) -> torch.Tensor | None: - return None if self.is_embed is None else self.is_embed.cumsum(dim=0) + def embeds_cumsum(self) -> list[int] | None: + # python list so python indexing avoids torch C++ overhead/conversions/deallocs + return None if self.is_embed is None else self.is_embed.cumsum(dim=0).tolist() def get_num_embeds(self) -> int: if self.embeds_cumsum is None: return self.length - return int(self.embeds_cumsum[-1]) + return self.embeds_cumsum[-1] if self.embeds_cumsum else 0 def get_embeds_indices_in_range( self, start_idx: int, end_idx: int @@ -170,10 +171,8 @@ class PlaceholderRange: if self.embeds_cumsum is None: return start_idx, end_idx - embeds_start_idx = ( - int(self.embeds_cumsum[start_idx - 1]) if start_idx > 0 else 0 - ) - embeds_end_idx = int(self.embeds_cumsum[end_idx - 1]) + embeds_start_idx = self.embeds_cumsum[start_idx - 1] if start_idx > 0 else 0 + embeds_end_idx = self.embeds_cumsum[end_idx - 1] if end_idx > 0 else 0 return embeds_start_idx, embeds_end_idx