support speculative execution for openai API (#48)

Co-authored-by: Ying Sheng <sqy1415@gmail.com>
2024-01-25 03:57:06 -06:00
parent 93414c8238
commit 23950056f0
10 changed files with 178 additions and 12 deletions
--- a/python/sglang/api.py
+++ b/python/sglang/api.py
@@ -20,8 +20,16 @@ from sglang.lang.ir import (
 )


-def function(func: Callable):
-    return SglFunction(func)
+def function(
+    func: Optional[Callable] = None, api_num_spec_tokens: Optional[int] = None
+):
+    if func:
+        return SglFunction(func, api_num_spec_tokens=api_num_spec_tokens)
+
+    def decorator(func):
+        return SglFunction(func, api_num_spec_tokens=api_num_spec_tokens)
+
+    return decorator


 def Runtime(*args, **kwargs):