fix: add tp_size for sglang_wrapper

Author: ChenZiHong-Gavin

graphgen/models/llm/local/sglang_wrapper.py

@@ -33,9 +33,10 @@ def __init__(
         self.temperature = temperature
         self.top_p = top_p
         self.topk = topk
+        self.tp_size = int(tp_size)
 
         # Initialise the offline engine
-        self.engine = sgl.Engine(model_path=self.model_path, tp_size=int(tp_size))
+        self.engine = sgl.Engine(model_path=self.model_path, tp_size=tp_size)
 
         # Keep helpers for streaming
         self.async_stream_and_merge = async_stream_and_merge
@@ -146,4 +147,6 @@ def shutdown(self) -> None:
     def restart(self) -> None:
         """Restart the SGLang engine."""
         self.shutdown()
-        self.engine = self.engine.__class__(model_path=self.model_path)
+        self.engine = self.engine.__class__(
+            model_path=self.model_path, tp_size=self.tp_size
+        )