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, 06:00 UTC
[Misc] refactor neuron_multimodal and profiling (#19397)
Signed-off-by: reidliu41 <reid201711@gmail.com> Co-authored-by: reidliu41 <reid201711@gmail.com>
Reid committed
Jun 10, 2025 at 14:12 UTC
6b1391ca7ece0350ecbe07e91f2c91d901c7f525
2 files changed
+11
-3
examples/offline_inference/neuron_multimodal.py
+5
-1
@@ -64,7 +64,7 @@ def print_outputs(outputs):
64
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
65
66
67
-if __name__ == "__main__":
67
+def main():
68
assert (
69
len(PROMPTS) == len(IMAGES) == len(SAMPLING_PARAMS)
70
), f"""Text, image prompts and sampling parameters should have the
@@ -104,3 +104,7 @@ if __name__ == "__main__":
104
# test batch-size = 4
105
outputs = llm.generate(batched_inputs, batched_sample_params)
106
print_outputs(outputs)
107
+
108
+
109
+if __name__ == "__main__":
110
+ main()
examples/offline_inference/profiling_tpu/profiling.py
+6
-2
@@ -70,7 +70,7 @@ def main(args: argparse.Namespace):
70
return
71
72
73
-if __name__ == "__main__":
73
+def parse_args():
74
parser = FlexibleArgumentParser(
75
description="Benchmark the latency of processing a single batch of "
76
"requests till completion."
@@ -102,5 +102,9 @@ if __name__ == "__main__":
102
)
103
104
parser = EngineArgs.add_cli_args(parser)
105
- args = parser.parse_args()
105
+ return parser.parse_args()
106
+
107
+
108
+if __name__ == "__main__":
109
+ args = parse_args()
110
main(args)