Upload serve_vllm.py with huggingface_hub

serve_vllm.py

@@ -0,0 +1,31 @@
+import os, sys, subprocess, argparse
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model', required=True)
+    parser.add_argument('--tensor_parallel_size', type=int, default=1)
+    parser.add_argument('--port', type=int, default=8000)
+    parser.add_argument('--max_model_len', type=int, default=8192)
+    parser.add_argument('--max_num_seqs', type=int, default=256)
+    parser.add_argument('--dtype', default='bfloat16')
+    parser.add_argument('--num_scheduler_steps', type=int, default=10)
+    parser.add_argument('--enable_chunked_prefill', action='store_true', default=False)
+    args = parser.parse_args()
+
+    os.environ.setdefault('VLLM_USE_TRITON_FLASH_ATTN', '0')
+    os.environ.setdefault('NCCL_MIN_NCHANNELS', '112')
+    os.environ.setdefault('PYTORCH_TUNABLEOP_ENABLED', '1')
+    os.environ.setdefault('TORCHINDUCTOR_MAX_AUTOTUNE', '1')
+
+    cmd = [sys.executable, '-m', 'vllm.entrypoints.openai.api_server',
+           '--model', args.model, '--tensor-parallel-size', str(args.tensor_parallel_size),
+           '--port', str(args.port), '--max-model-len', str(args.max_model_len),
+           '--max-num-seqs', str(args.max_num_seqs), '--dtype', args.dtype,
+           '--num-scheduler-steps', str(args.num_scheduler_steps)]
+    if not args.enable_chunked_prefill:
+        cmd.append('--enable-chunked-prefill=False')
+
+    print('Launching vLLM:', ' '.join(cmd))
+    subprocess.run(cmd)
+
+if __name__ == "__main__": main()