A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 01:00 UTC
[Core] Reduce mm scheduler, get_num_embed overhead (#40143)
Signed-off-by: milesial <milesial@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
milesial committed
Apr 17, 2026 at 20:25 UTC
b0755523dce2570d74795be2e7462f2ce7d83bfd
2 files changed
+9
-13
tests/multimodal/test_inputs.py
+3
-6
@@ -26,11 +26,8 @@ def test_placeholder_range_get_num_embeds(is_embed, expected):
26
"is_embed,expected",
27
[
28
(None, None),
29
- (
30
- torch.tensor([False, True, False, True, True]),
31
- torch.tensor([0, 1, 1, 2, 3]),
32
- ),
33
- (torch.tensor([True, True, True]), torch.tensor([1, 2, 3])),
29
+ (torch.tensor([False, True, False, True, True]), [0, 1, 1, 2, 3]),
30
+ (torch.tensor([True, True, True]), [1, 2, 3]),
31
],
32
)
33
def test_placeholder_range_embeds_cumsum(is_embed, expected):
@@ -41,6 +38,6 @@ def test_placeholder_range_embeds_cumsum(is_embed, expected):
38
assert pr.embeds_cumsum is None
39
return
40
44
- assert torch.equal(pr.embeds_cumsum, expected)
41
+ assert pr.embeds_cumsum == expected
42
# cached_property should return the same object on repeated access
43
assert pr.embeds_cumsum is pr.embeds_cumsum
vllm/multimodal/inputs.py
+6
-7
@@ -145,14 +145,15 @@ class PlaceholderRange:
145
"""
146
147
@cached_property
148
- def embeds_cumsum(self) -> torch.Tensor | None:
149
- return None if self.is_embed is None else self.is_embed.cumsum(dim=0)
148
+ def embeds_cumsum(self) -> list[int] | None:
149
+ # python list so python indexing avoids torch C++ overhead/conversions/deallocs
150
+ return None if self.is_embed is None else self.is_embed.cumsum(dim=0).tolist()
151
152
def get_num_embeds(self) -> int:
153
if self.embeds_cumsum is None:
154
return self.length
155
155
- return int(self.embeds_cumsum[-1])
156
+ return self.embeds_cumsum[-1] if self.embeds_cumsum else 0
157
158
def get_embeds_indices_in_range(
159
self, start_idx: int, end_idx: int
@@ -170,10 +171,8 @@ class PlaceholderRange:
171
if self.embeds_cumsum is None:
172
return start_idx, end_idx
173
173
- embeds_start_idx = (
174
- int(self.embeds_cumsum[start_idx - 1]) if start_idx > 0 else 0
175
- )
176
- embeds_end_idx = int(self.embeds_cumsum[end_idx - 1])
174
+ embeds_start_idx = self.embeds_cumsum[start_idx - 1] if start_idx > 0 else 0
175
+ embeds_end_idx = self.embeds_cumsum[end_idx - 1] if end_idx > 0 else 0
176
177
return embeds_start_idx, embeds_end_idx
178