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 7, 02:45 UTC
refine `vllm bench throughput --backend hf` (#35971)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Kunshang Ji committed
Mar 7, 2026 at 10:10 UTC
7eb524e64c4533a5e24909873bb926109f3a4ac7
1 file changed
+16
-3
vllm/benchmarks/throughput.py
+16
-3
@@ -38,6 +38,7 @@ from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
38
from vllm.inputs import TextPrompt, TokensPrompt
39
from vllm.lora.request import LoRARequest
40
from vllm.outputs import RequestOutput
41
+from vllm.platforms import current_platform
42
from vllm.sampling_params import BeamSearchParams
43
from vllm.tokenizers import TokenizerLike, get_tokenizer
44
from vllm.utils.async_utils import merge_async_iterators
@@ -256,17 +257,21 @@ def run_hf(
257
max_batch_size: int,
258
trust_remote_code: bool,
259
disable_detokenize: bool = False,
260
+ dtype: torch.dtype | None = torch.float16,
261
+ enable_torch_compile: bool = False,
262
) -> float:
263
assert isinstance(tokenizer, PreTrainedTokenizerBase), (
264
"the hf backend only supports HF tokenizers"
265
)
266
llm = AutoModelForCausalLM.from_pretrained(
264
- model, dtype=torch.float16, trust_remote_code=trust_remote_code
267
+ model, dtype=dtype, trust_remote_code=trust_remote_code
268
)
269
if llm.config.model_type == "llama":
270
# To enable padding in the HF backend.
271
tokenizer.pad_token = tokenizer.eos_token
269
- llm = llm.cuda()
272
+ llm = llm.to(current_platform.device_type)
273
+ if enable_torch_compile:
274
+ llm = torch.compile(llm)
275
276
pbar = tqdm(total=len(requests))
277
start = time.perf_counter()
@@ -295,7 +300,7 @@ def run_hf(
300
# Generate the sequences.
301
input_ids = tokenizer(batch, return_tensors="pt", padding=True).input_ids
302
llm_outputs = llm.generate(
298
- input_ids=input_ids.cuda(),
303
+ input_ids=input_ids.to(current_platform.device_type),
304
do_sample=True,
305
num_return_sequences=n,
306
temperature=1.0,
@@ -733,6 +738,12 @@ def add_cli_args(parser: argparse.ArgumentParser):
738
default=None,
739
help="Maximum batch size for HF backend.",
740
)
741
+ parser.add_argument(
742
+ "--hf-enable-torch-compile",
743
+ action="store_true",
744
+ default=False,
745
+ help="Enable Torch compile for HF backend.",
746
+ )
747
parser.add_argument(
748
"--output-json",
749
type=str,
@@ -884,6 +895,8 @@ def main(args: argparse.Namespace):
895
args.hf_max_batch_size,
896
args.trust_remote_code,
897
args.disable_detokenize,
898
+ dtype=args.dtype,
899
+ enable_torch_compile=args.hf_enable_torch_compile,
900
)
901
elif args.backend == "vllm-chat":
902
elapsed_time, request_outputs = run_vllm_chat(