Upload serve_vllm.py

serve_vllm.py

@@ -0,0 +1,32 @@
+"""vLLM OpenAI-compatible API server for AMD ROCm MI300X."""
+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()