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, 21:00 UTC
[Perf] Remove redundant device copies for CPU-only pooling token IDs, 48.9% E2E throughput improvement (#38139)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Mar 29, 2026 at 14:12 UTC
995dea1354cbbb378f837b6897af420fd2f75322
8 files changed
+86
-17
tests/v1/worker/test_gpu_input_batch.py
+47
-2
@@ -380,7 +380,7 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
380
_compare_objs(input_batch, ref_input_batch)
381
382
383
-def _construct_pooling_request(req_id_suffix: int):
383
+def _construct_pooling_request(req_id_suffix: int, pooling_params=None):
384
from vllm.pooling_params import PoolingParams
385
386
prompt_token_ids = [
@@ -391,7 +391,7 @@ def _construct_pooling_request(req_id_suffix: int):
391
req_id=f"pool_req_{req_id_suffix}",
392
prompt_token_ids=prompt_token_ids,
393
sampling_params=None,
394
- pooling_params=PoolingParams(task="classify"),
394
+ pooling_params=pooling_params or PoolingParams(task="classify"),
395
mm_features=[],
396
block_ids=([],),
397
generator=None,
@@ -440,3 +440,48 @@ def test_pooling_prompt_lens_not_aliased(device: str):
440
"mutations to num_prompt_tokens_cpu_tensor corrupted prompt_lens. "
441
f"Expected {prompt_lens_snapshot}, got {metadata.prompt_lens}"
442
)
443
+
444
+
445
+@pytest.mark.parametrize(
446
+ ("pooling_params", "expect_device_prompt_token_ids", "expect_cpu_prompt_token_ids"),
447
+ [
448
+ ({"task": "classify"}, False, False),
449
+ ({"task": "classify", "requires_token_ids": True}, True, True),
450
+ ],
451
+)
452
+def test_pooling_metadata_token_id_buffers(
453
+ pooling_params: dict[str, object],
454
+ expect_device_prompt_token_ids: bool,
455
+ expect_cpu_prompt_token_ids: bool,
456
+):
457
+ from vllm.pooling_params import PoolingParams
458
+
459
+ input_batch = InputBatch(
460
+ max_num_reqs=1,
461
+ max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
462
+ max_num_batched_tokens=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
463
+ device=torch.device("cpu"),
464
+ pin_memory=False,
465
+ vocab_size=VOCAB_SIZE,
466
+ block_sizes=[16],
467
+ kernel_block_sizes=[16],
468
+ is_pooling_model=True,
469
+ )
470
+ req = _construct_pooling_request(0, PoolingParams(**pooling_params))
471
+ input_batch.add_request(req)
472
+ input_batch.refresh_metadata()
473
+
474
+ metadata = input_batch.get_pooling_metadata()
475
+ if expect_device_prompt_token_ids:
476
+ assert input_batch.sampling_metadata.prompt_token_ids is not None
477
+ assert metadata.prompt_token_ids is not None
478
+ assert metadata.get_prompt_token_ids()[0].tolist() == req.prompt_token_ids
479
+ else:
480
+ assert input_batch.sampling_metadata.prompt_token_ids is None
481
+ assert metadata.prompt_token_ids is None
482
+
483
+ if expect_cpu_prompt_token_ids:
484
+ assert metadata.prompt_token_ids_cpu is not None
485
+ assert metadata.get_prompt_token_ids_cpu()[0].tolist() == req.prompt_token_ids
486
+ else:
487
+ assert metadata.prompt_token_ids_cpu is None
vllm/model_executor/layers/pooler/common.py
+1
-1
@@ -18,7 +18,7 @@ ActivationFn = Callable[[_T], _T]
18
@dataclass(frozen=True)
19
class PoolingParamsUpdate:
20
requires_token_ids: bool = False
21
- """Set this flag to enable `get_prompt_token_ids` for your pooler."""
21
+ """Set this flag to enable prompt token IDs for your pooler."""
22
23
def __or__(self, other: "PoolingParamsUpdate") -> "PoolingParamsUpdate":
24
return PoolingParamsUpdate(
vllm/model_executor/layers/pooler/special.py
+6
-4
@@ -146,17 +146,19 @@ class BOSEOSFilter(Pooler):
146
) -> PoolerOutput:
147
pooled_outputs = self.pooler(hidden_states, pooling_metadata)
148
assert isinstance(pooled_outputs, list)
149
+ prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu()
150
150
- for i, prompt_len in enumerate(pooling_metadata.prompt_lens):
151
+ for i, (prompt_len, token_ids) in enumerate(
152
+ zip(pooling_metadata.prompt_lens, prompt_token_ids)
153
+ ):
154
pooled_data = pooled_outputs[i]
155
assert (
156
isinstance(pooled_data, torch.Tensor)
157
and pooled_data.shape[0] == prompt_len
158
)
156
- token_ids = pooling_metadata.prompt_token_ids[i, :prompt_len]
157
- if token_ids[0] == self.bos_token_id:
159
+ if int(token_ids[0]) == self.bos_token_id:
160
pooled_data = pooled_data[1:]
159
- if token_ids[-1] == self.eos_token_id:
161
+ if int(token_ids[-1]) == self.eos_token_id:
162
pooled_data = pooled_data[:-1]
163
pooled_outputs[i] = pooled_data.squeeze(-1)
164
vllm/model_executor/models/bert.py
+5
-4
@@ -638,25 +638,26 @@ class SPLADESparsePooler(Pooler):
638
lens: list[int] = lens_tensor.tolist()
639
B: int = len(lens)
640
641
- token_ids = pooling_metadata.prompt_token_ids
641
+ prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu()
642
offset = 0
643
pooled_list: list[torch.Tensor] = []
644
645
for i in range(B):
646
L = int(lens[i])
647
hs = hidden_states[offset : offset + L]
648
+ token_ids = prompt_token_ids[i]
649
650
start_idx = 0
651
end_idx = L
651
- if self.remove_cls_sep and token_ids is not None:
652
+ if self.remove_cls_sep:
653
if (
654
self.cls_token_id is not None
654
- and token_ids[i, 0].item() == self.cls_token_id
655
+ and int(token_ids[0]) == self.cls_token_id
656
):
657
start_idx = 1
658
if (
659
self.sep_token_id is not None
659
- and token_ids[i, L - 1].item() == self.sep_token_id
660
+ and int(token_ids[L - 1]) == self.sep_token_id
661
):
662
end_idx = max(start_idx, L - 1)
663
vllm/model_executor/models/gritlm.py
+3
-2
@@ -156,10 +156,11 @@ class GritLMMeanPool(SequencePoolingMethod):
156
pooling_metadata: PoolingMetadata,
157
) -> SequencePoolingMethodOutput:
158
prompt_lens = pooling_metadata.prompt_lens
159
+ prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu()
160
instr_lens = torch.tensor(
161
[
161
- self._get_instruction_len(token_ids.cpu().numpy())
162
- for token_ids in pooling_metadata.get_prompt_token_ids()
162
+ self._get_instruction_len(token_ids.numpy())
163
+ for token_ids in prompt_token_ids
164
],
165
device="cpu",
166
)
vllm/v1/pool/metadata.py
+11
-1
@@ -50,7 +50,8 @@ class PoolingMetadata:
50
"""Tensors for pooling."""
51
52
prompt_lens: torch.Tensor # CPU Tensor
53
- prompt_token_ids: torch.Tensor | None
53
+ prompt_token_ids: torch.Tensor | None # Model-device tensor
54
+ prompt_token_ids_cpu: torch.Tensor | None # CPU tensor
55
pooling_params: list[PoolingParams]
56
pooling_states: list[PoolingStates]
57
pooling_cursor: PoolingCursor | None = None
@@ -73,6 +74,9 @@ class PoolingMetadata:
74
prompt_token_ids=None
75
if self.prompt_token_ids is None
76
else self.prompt_token_ids[indices],
77
+ prompt_token_ids_cpu=None
78
+ if self.prompt_token_ids_cpu is None
79
+ else self.prompt_token_ids_cpu[indices],
80
pooling_params=self.pooling_params[indices],
81
pooling_states=self.pooling_states[indices],
82
pooling_cursor=None
@@ -85,7 +89,13 @@ class PoolingMetadata:
89
assert prompt_token_ids is not None, (
90
"Please set `requires_token_ids=True` in `get_pooling_updates`"
91
)
92
+ return [prompt_token_ids[i, :num] for i, num in enumerate(self.prompt_lens)]
93
94
+ def get_prompt_token_ids_cpu(self) -> list[torch.Tensor]:
95
+ prompt_token_ids = self.prompt_token_ids_cpu
96
+ assert prompt_token_ids is not None, (
97
+ "Please set `requires_token_ids=True` in `get_pooling_updates`"
98
+ )
99
return [prompt_token_ids[i, :num] for i, num in enumerate(self.prompt_lens)]
100
101
def get_pooling_cursor(self) -> PoolingCursor:
vllm/v1/worker/gpu_input_batch.py
+12
-3
@@ -833,8 +833,13 @@ class InputBatch:
833
# step pooling during the sampling/pooling process.
834
# Hence copy these tensors only when there are requests which
835
# need penalties/step_pooler to be applied.
836
+ prompt_token_ids_cpu = (
837
+ self._make_prompt_token_ids_cpu_tensor() if needs_prompt_token_ids else None
838
+ )
839
prompt_token_ids = (
837
- self._make_prompt_token_ids_tensor() if needs_prompt_token_ids else None
840
+ prompt_token_ids_cpu.to(device=self.device, non_blocking=True)
841
+ if prompt_token_ids_cpu is not None
842
+ else None
843
)
844
845
# Only set output_token_ids if required by the current requests'
@@ -891,15 +896,19 @@ class InputBatch:
896
def get_pooling_metadata(self) -> PoolingMetadata:
897
pooling_params = self.get_pooling_params()
898
pooling_states = self.get_pooling_states()
899
+ prompt_token_ids_cpu = None
900
+ if any(p.requires_token_ids for p in pooling_params):
901
+ prompt_token_ids_cpu = self._make_prompt_token_ids_cpu_tensor()
902
903
return PoolingMetadata(
904
prompt_lens=self.num_prompt_tokens_cpu_tensor[: self.num_reqs].clone(),
905
prompt_token_ids=self.sampling_metadata.prompt_token_ids,
906
+ prompt_token_ids_cpu=prompt_token_ids_cpu,
907
pooling_params=pooling_params,
908
pooling_states=pooling_states,
909
)
910
902
- def _make_prompt_token_ids_tensor(self) -> torch.Tensor:
911
+ def _make_prompt_token_ids_cpu_tensor(self) -> torch.Tensor:
912
num_reqs = self.num_reqs
913
max_prompt_len = self.num_prompt_tokens[:num_reqs].max()
914
prompt_token_ids_cpu_tensor = torch.empty(
@@ -914,7 +923,7 @@ class InputBatch:
923
# token_id of this value.
924
for i in range(num_reqs):
925
prompt_token_ids[i, self.num_prompt_tokens[i] :] = self.vocab_size
917
- return prompt_token_ids_cpu_tensor.to(device=self.device, non_blocking=True)
926
+ return prompt_token_ids_cpu_tensor
927
928
def make_lora_inputs(
929
self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray
vllm/v1/worker/gpu_model_runner.py
+1
@@ -5653,6 +5653,7 @@ class GPUModelRunner(
5653
dummy_metadata = PoolingMetadata(
5654
prompt_lens=dummy_prompt_lens,
5655
prompt_token_ids=dummy_token_ids,
5656
+ prompt_token_ids_cpu=dummy_token_ids.cpu(),
5657
pooling_params=[dummy_pooling_params] * num_reqs,
5658
pooling_states=[PoolingStates() for i in range(num_reqs)],
5659
)