support speculative execution for openai API (#48)

Co-authored-by: Ying Sheng <sqy1415@gmail.com>

test/lang/test_openai_spec.py

@@ -0,0 +1,68 @@
+from sglang import OpenAI, function, gen, set_default_backend
+
+
+@function()
+def gen_character_default(s):
+    s += "Construct a character within the following format:\n"
+    s += "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\nWelcome.\n"
+    s += "\nPlease generate new Name, Birthday and Job.\n"
+    s += "Name:" + gen("name", stop="\n") + "\nBirthday:" + gen("birthday", stop="\n")
+    s += "\nJob:" + gen("job", stop="\n") + "\nWelcome.\n"
+
+
+@function(api_num_spec_tokens=512)
+def gen_character_spec(s):
+    s += "Construct a character within the following format:\n"
+    s += "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\nWelcome.\n"
+    s += "\nPlease generate new Name, Birthday and Job.\n"
+    s += "Name:" + gen("name", stop="\n") + "\nBirthday:" + gen("birthday", stop="\n")
+    s += "\nJob:" + gen("job", stop="\n") + "\nWelcome.\n"
+
+
+@function(api_num_spec_tokens=512)
+def gen_character_no_stop(s):
+    s += "Construct a character within the following format:\n"
+    s += "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\nWelcome.\n"
+    s += "\nPlease generate new Name, Birthday and Job.\n"
+    s += "Name:" + gen("name") + "\nBirthday:" + gen("birthday")
+    s += "\nJob:" + gen("job") + "\nWelcome.\n"
+
+
+@function(api_num_spec_tokens=512)
+def gen_character_multi_stop(s):
+    s += "Construct a character within the following format:\n"
+    s += (
+        "Name: Steve Jobs.###Birthday: February 24, 1955.###Job: Apple CEO.\nWelcome.\n"
+    )
+    s += "\nPlease generate new Name, Birthday and Job.\n"
+    s += "Name:" + gen("name", stop=["\n", "###"])
+    s += "###Birthday:" + gen("birthday", stop=["\n", "###"])
+    s += "###Job:" + gen("job", stop=["\n", "###"]) + "\nWelcome.\n"
+
+
+set_default_backend(OpenAI("gpt-3.5-turbo-instruct"))
+
+state = gen_character_default.run()
+print(state.text())
+
+print("=" * 60)
+
+state = gen_character_no_stop.run()
+
+print("name###", state["name"])
+print("birthday###:", state["birthday"])
+print("job###", state["job"])
+
+print("=" * 60)
+
+state = gen_character_multi_stop.run()
+print(state.text())
+
+print("=" * 60)
+
+state = gen_character_spec.run()
+print(state.text())
+
+print("name###", state["name"])
+print("birthday###", state["birthday"])
+print("job###", state["job"])

test/srt/test_fast_forward.py

@@ -1,7 +1,6 @@
 import argparse
 from enum import Enum
 
-import sglang as sgl
 from pydantic import BaseModel, constr
 from sglang.srt.constrained.json_schema import build_regex_from_object
 from sglang.test.test_utils import (
@@ -9,6 +8,8 @@ from sglang.test.test_utils import (
     select_sglang_backend,
 )
 
+import sglang as sgl
+
 IP_REGEX = r"((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)"
 
 ip_fast_forward = (

test/srt/test_robust.py

@@ -2,13 +2,14 @@ import argparse
 import random
 import string
 
-import sglang as sgl
 from sglang.test.test_utils import (
     add_common_sglang_args_and_parse,
     select_sglang_backend,
 )
 from vllm.transformers_utils.tokenizer import get_tokenizer
 
+import sglang as sgl
+
 TOKENIZER = None
 RANDOM_PREFILL_LEN = None
 RANDOM_DECODE_LEN = None