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 15, 22:15 UTC
[Bugfix] V1 Memory Profiling: V0 Sampler Integration without Rejection Sampler (#13594)
Signed-off-by: Jennifer Zhao <7443418+JenZhao@users.noreply.github.com> Co-authored-by: Roger Wang <ywang@roblox.com>
Jennifer Zhao committed
Feb 22, 2025 at 00:08 UTC
da31b5333ec198343a79cd8b5b198ac76896de51
1 file changed
+26
-2
vllm/v1/worker/gpu_model_runner.py
+26
-2
@@ -31,6 +31,7 @@ from vllm.v1.engine.mm_input_cache import MMInputCacheClient
31
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
32
KVCacheSpec)
33
from vllm.v1.outputs import LogprobsTensors, ModelRunnerOutput
34
+from vllm.v1.sample.metadata import SamplingMetadata
35
from vllm.v1.sample.rejection_sampler import INVALID_TOKEN_ID
36
from vllm.v1.spec_decode.ngram_proposer import NgramProposer
37
from vllm.v1.utils import bind_kv_cache
@@ -1305,11 +1306,34 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1306
if get_pp_group().is_last_rank:
1307
hidden_states = hidden_states[logit_indices]
1308
logits = self.model.compute_logits(hidden_states, None)
1308
- # TODO(woosuk): Consider the memory usage of the sampler.
1309
+ dummy_tensors = lambda v: torch.full(
1310
+ (num_reqs, ), v, device=self.device)
1311
+ dummy_metadata = SamplingMetadata(
1312
+ temperature=dummy_tensors(0.5),
1313
+ all_greedy=False,
1314
+ all_random=False,
1315
+ spec_token_ids=None,
1316
+ top_p=dummy_tensors(0.9),
1317
+ top_k=dummy_tensors(logits.size(1) - 1),
1318
+ min_p=None,
1319
+ generators={},
1320
+ max_num_logprobs=None,
1321
+ no_penalties=True,
1322
+ prompt_token_ids=torch.ones_like(logits, dtype=torch.int64),
1323
+ frequency_penalties=dummy_tensors(0.1),
1324
+ presence_penalties=dummy_tensors(0.1),
1325
+ repetition_penalties=dummy_tensors(0.1),
1326
+ output_token_ids=[[] for _ in range(num_reqs)],
1327
+ min_tokens={},
1328
+ logit_bias=[None for _ in range(num_reqs)])
1329
+ sampler_output = self.model.sample(
1330
+ logits=logits, sampling_metadata=dummy_metadata)
1331
else:
1332
logits = None
1333
+ sampler_output = None
1334
+ dummy_metadata = None
1335
torch.cuda.synchronize()
1312
- del hidden_states, logits
1336
+ del hidden_states, logits, sampler_output, dummy_metadata
1337
self.encoder_cache.clear()
1338
gc.collect()
1339