support spec_tokens=512

examples/quick_start/openai_example_speculative.py

@@ -1,5 +1,6 @@
 from sglang import function, gen, set_default_backend, OpenAI
-@function
+
+@function(spec_tokens=512)
 def example(s):
   s += "Construct a character. Here is an example:\n"
   s += "Name: Steve Jobs. Birthday: February 24, 1955. Job: Apple CEO.\n"

python/sglang/api.py

@@ -19,10 +19,12 @@
     SglSelect,
 )
 
-
-def function(func: Callable):
-    return SglFunction(func)
-
+def function(func=None,**bind_arguments):
+    if func:
+        return SglFunction(func, bind_arguments)
+    def decorator(func):
+        return SglFunction(func, bind_arguments)
+    return decorator
 
 def Runtime(*args, **kwargs):
     # Avoid importing unnecessary dependency

python/sglang/lang/interpreter.py

@@ -49,10 +49,11 @@ def run_program(
     program, backend, func_args, func_kwargs, default_sampling_para, stream, sync=False
 ):
     assert backend is not None, "Please specify a backend"
-    func_kwargs.update(program.bind_arguments)
+    # func_kwargs.update(program.bind_arguments)
     stream_executor = StreamExecutor(
         backend, func_kwargs, default_sampling_para, chat_template=None, stream=stream
     )
+    stream_executor.bind_arguments = program.bind_arguments
     state = ProgramState(stream_executor)
 
     if stream:
@@ -170,7 +171,8 @@ def __init__(
 
         # For speculative execution
         self.last_comp = "" # The last completion message returned 
-        self.speculate_text_prefix = "" # The speculated text prefix 
+        self.speculate_text_prefix = "" # The speculated text prefix
+        self.bind_arguments = {}
 
         # For chat
         self.messages_ = []  # The messages in the OpenAI API format
@@ -361,6 +363,11 @@ def _execute_gen(self, expr: SglGen):
                 self.variables[self.last_name] = last_comp
                 meta_info = {}
             else:
+                if self.last_comp == "" and 'spec_tokens' in self.bind_arguments:
+                    spec_tokens_value = self.bind_arguments['spec_tokens']
+                    sampling_params.max_new_tokens = spec_tokens_value
+                    sampling_params.stop = None
+
                 comp, meta_info = self.backend.generate(
                     self, sampling_params=sampling_params
                 )