Repository: sgl-project/sglang
Pull request #3612: Feat/support code completion
Author: login: woodx9
Created: 2025-02-16T14:35:00Z | Merged: 2025-03-17T01:26:20Z

## Motivation

<!-- Thank you for your contribution! We appreciate it. The following guidelines will help improve your pull request and facilitate feedback. If anything is unclear, don't hesitate to submit your pull request and ask the maintainers for assistance. -->

support code completion in open api v1/completion. discussion is held in [issue](https://github.com/sgl-project/sglang/issues/3429)
this feature has been test by myself, I will add unit test after the feature is reviewed, thanks!
@shuaills cc

## Modifications

1. add completion template in launch server paramter
2. add completion template class like chat template
3. adapt the prompt for  completion template before it's processed by engine


- [x] Format your code according to the [Code Formatting with Pre-Commit](https://docs.sglang.ai/references/contribution_guide.html#code-formatting-with-pre-commit).
- [x] Add unit tests as outlined in the [Running Unit Tests](https://docs.sglang.ai/references/contribution_guide.html#running-unit-tests-adding-to-ci).
- [ ] Update documentation / docstrings / example tutorials as needed, according to [Writing Documentation](https://docs.sglang.ai/references/contribution_guide.html#writing-documentation-running-docs-ci).
- [ ] Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to [Benchmark and Profiling](https://docs.sglang.ai/references/benchmark_and_profiling.html) and [Accuracy Results](https://docs.sglang.ai/references/accuracy_evaluation.html).
- [ ] For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
- [ ] Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.



Diff:
diff --git a/python/sglang/srt/code_completion_parser.py b/python/sglang/srt/code_completion_parser.py
new file mode 100644
index 000000000000..94a98b0fdc48
--- /dev/null
+++ b/python/sglang/srt/code_completion_parser.py
@@ -0,0 +1,174 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Completion templates."""
+
+
+import dataclasses
+import json
+import logging
+import os
+from enum import auto
+
+from sglang.srt.openai_api.protocol import ChatCompletionRequest
+
+logger = logging.getLogger(__name__)
+completion_template_name = None
+
+
+class FimPosition:
+    """Postion of fim middle token."""
+
+    MIDDLE = auto()
+    END = auto()
+
+
+@dataclasses.dataclass
+class CompletionTemplate:
+    """A class that manages completion prompt templates. only for code completion currently."""
+
+    # The name of this template
+    name: str
+
+    # the fim begin token
+    fim_begin_token: str
+
+    # The fim middle token
+    fim_middle_token: str
+
+    # The fim end token
+    fim_end_token: str
+
+    # The position of the fim middle token
+    fim_position: FimPosition
+
+
+# A global registry for all completion templates
+completion_templates: dict[str, CompletionTemplate] = {}
+
+
+def load_completion_template_for_openai_api(completion_template_arg):
+    global completion_template_name
+
+    logger.info(
+        f"Use completion template for the OpenAI-compatible API server: {completion_template_arg}"
+    )
+
+    if not completion_template_exists(completion_template_arg):
+        if not os.path.exists(completion_template_arg):
+            raise RuntimeError(
+                f"Completion template {completion_template_arg} is not a built-in template name "
+                "or a valid completion template file path."
+            )
+
+        assert completion_template_arg.endswith(
+            ".json"
+        ), "unrecognized format of completion template file"
+        with open(completion_template_arg, "r") as filep:
+            template = json.load(filep)
+            try:
+                fim_position = FimPosition[template["fim_position"]]
+            except KeyError:
+                raise ValueError(
+                    f"Unknown fim position: {template['fim_position']}"
+                ) from None
+            register_completion_template(
+                CompletionTemplate(
+                    name=template["name"],
+                    fim_begin_token=template["fim_begin_token"],
+                    fim_middle_token=template["fim_middle_token"],
+                    fim_end_token=template["fim_end_token"],
+                    fim_position=fim_position,
+                ),
+                override=True,
+            )
+        completion_template_name = template["name"]
+    else:
+        completion_template_name = completion_template_arg
+
+
+def register_completion_template(template: CompletionTemplate, override: bool = False):
+    """Register a new completion template."""
+    if not override:
+        assert (
+            template.name not in completion_templates
+        ), f"{template.name} has been registered."
+
+    completion_templates[template.name] = template
+
+
+def completion_template_exists(template_name: str) -> bool:
+    return template_name in completion_templates
+
+
+def is_completion_template_defined() -> bool:
+    global completion_template_name
+    return completion_template_name != None
+
+
+def generate_completion_prompt_from_request(request: ChatCompletionRequest) -> str:
+    global completion_template_name
+    if request.suffix == "":
+        return request.prompt
+
+    return generate_completion_prompt(
+        request.prompt, request.suffix, completion_template_name
+    )
+
+
+def generate_completion_prompt(prompt: str, suffix: str, template_name: str) -> str:
+
+    completion_template = completion_templates[template_name]
+    fim_begin_token = completion_template.fim_begin_token
+    fim_middle_token = completion_template.fim_middle_token
+    fim_end_token = completion_template.fim_end_token
+    fim_position = completion_template.fim_position
+
+    if fim_position == FimPosition.MIDDLE:
+        prompt = f"{fim_begin_token}{prompt}{fim_middle_token}{suffix}{fim_end_token}"
+    elif fim_position == FimPosition.END:
+        prompt = f"{fim_begin_token}{prompt}{fim_end_token}{suffix}{fim_middle_token}"
+
+    return prompt
+
+
+register_completion_template(
+    CompletionTemplate(
+        name="deepseek_coder",
+        fim_begin_token="<｜fim▁begin｜>",
+        fim_middle_token="<｜fim▁hole｜>",
+        fim_end_token="<｜fim▁end｜>",
+        fim_position=FimPosition.MIDDLE,
+    )
+)
+
+
+register_completion_template(
+    CompletionTemplate(
+        name="star_coder",
+        fim_begin_token="<fim_prefix>",
+        fim_middle_token="<fim_middle>",
+        fim_end_token="<fim_suffix>",
+        fim_position=FimPosition.END,
+    )
+)
+
+register_completion_template(
+    CompletionTemplate(
+        name="qwen_coder",
+        fim_begin_token="
",
+        fim_middle_token="
",
+        fim_end_token="
",
+        fim_position=FimPosition.END,
+    )
+)
diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py
index ec4ea515b4b2..8fc792b6c011 100644
--- a/python/sglang/srt/entrypoints/engine.py
+++ b/python/sglang/srt/entrypoints/engine.py
@@ -36,6 +36,7 @@
 import torch
 import uvloop
 
+from sglang.srt.code_completion_parser import load_completion_template_for_openai_api
 from sglang.srt.managers.data_parallel_controller import (
     run_data_parallel_controller_process,
 )
@@ -538,6 +539,9 @@ def _launch_subprocesses(
             tokenizer_manager, server_args.chat_template, server_args.model_path
         )
 
+    if server_args.completion_template:
+        load_completion_template_for_openai_api(server_args.completion_template)
+
     # Wait for the model to finish loading
     scheduler_infos = []
     for i in range(len(scheduler_pipe_readers)):
diff --git a/python/sglang/srt/openai_api/adapter.py b/python/sglang/srt/openai_api/adapter.py
index 2ac4e3ed85ae..ba4ca7b4ffc4 100644
--- a/python/sglang/srt/openai_api/adapter.py
+++ b/python/sglang/srt/openai_api/adapter.py
@@ -33,6 +33,10 @@
     # outlines.integrations.utils
     from outlines.integrations.utils import convert_json_schema_to_str
 
+from sglang.srt.code_completion_parser import (
+    generate_completion_prompt_from_request,
+    is_completion_template_defined,
+)
 from sglang.srt.conversation import (
     Conversation,
     SeparatorStyle,
@@ -504,7 +508,11 @@ def v1_generate_request(
                 "To compute logprobs of input prompt, please use the native /generate API."
             )
 
-        prompts.append(request.prompt)
+        prompt = request.prompt
+        if is_completion_template_defined():
+            prompt = generate_completion_prompt_from_request(request)
+        prompts.append(prompt)
+
         lora_paths.append(request.lora_path)
         if request.echo and request.logprobs:
             current_logprob_start_len = 0
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index d6cb878b943c..d31e6f000b7c 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -55,6 +55,7 @@ class ServerArgs:
     device: str = "cuda"
     served_model_name: Optional[str] = None
     chat_template: Optional[str] = None
+    completion_template: Optional[str] = None
     is_embedding: bool = False
     revision: Optional[str] = None
 
@@ -451,6 +452,12 @@ def add_cli_args(parser: argparse.ArgumentParser):
             default=ServerArgs.chat_template,
             help="The buliltin chat template name or the path of the chat template file. This is only used for OpenAI-compatible API server.",
         )
+        parser.add_argument(
+            "--completion-template",
+            type=str,
+            default=ServerArgs.completion_template,
+            help="The buliltin completion template name or the path of the completion template file. This is only used for OpenAI-compatible API server. only for code completion currently.",
+        )
         parser.add_argument(
             "--is-embedding",
             action="store_true",
diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py
index ecd1a39dd494..fde3805d2d44 100644
--- a/test/srt/run_suite.py
+++ b/test/srt/run_suite.py
@@ -70,6 +70,7 @@ class TestFile:
         TestFile("test_vision_chunked_prefill.py", 223),
         TestFile("test_vision_llm.py", 18.4),
         TestFile("test_vision_openai_server.py", 344),
+        TestFile("test_fim_completion.py", 120),
         TestFile("test_w8a8_quantization.py", 46),
         TestFile("test_eval_fp8_accuracy.py", 172),
     ],
diff --git a/test/srt/test_fim_completion.py b/test/srt/test_fim_completion.py
new file mode 100644
index 000000000000..132911e659bb
--- /dev/null
+++ b/test/srt/test_fim_completion.py
@@ -0,0 +1,71 @@
+import unittest
+
+import openai
+
+from sglang.srt.hf_transformers_utils import get_tokenizer
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    popen_launch_server,
+)
+
+
+class TestFimCompletion(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "deepseek-ai/deepseek-coder-1.3b-base"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.api_key = "sk-123456"
+        other_args = ["--completion-template", "deepseek_coder"]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            api_key=cls.api_key,
+            other_args=other_args,
+        )
+        cls.base_url += "/v1"
+        cls.tokenizer = get_tokenizer(cls.model)
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def run_fim_completion(self, number_of_completion):
+        client = openai.Client(api_key=self.api_key, base_url=self.base_url)
+        prompt = "function sum(a: number, b: number): number{\n"
+        suffix = "}"
+
+        prompt_input = self.tokenizer.encode(prompt) + self.tokenizer.encode(suffix)
+        num_prompt_tokens = len(prompt_input) + 2
+
+        response = client.completions.create(
+            model=self.model,
+            prompt=prompt,
+            suffix=suffix,
+            temperature=0.3,
+            max_tokens=32,
+            stream=False,
+            n=number_of_completion,
+        )
+
+        print(response)
+        print(len(response.choices))
+        assert len(response.choices) == number_of_completion
+        assert response.id
+        assert response.created
+        assert response.object == "text_completion"
+        assert (
+            response.usage.prompt_tokens == num_prompt_tokens
+        ), f"{response.usage.prompt_tokens} vs {num_prompt_tokens}"
+        assert response.usage.completion_tokens > 0
+        assert response.usage.total_tokens > 0
+
+    def test_fim_completion(self):
+        for number_of_completion in [1, 3]:
+            self.run_fim_completion(number_of_completion)
+
+
+if __name__ == "__main__":
+    unittest.main()


Reviews:
user: login: shuaills
body: Please add docs to https://docs.sglang.ai/backend/server_arguments.html and Advanced Features. For docs, you can have a look of [Structured Outputs](https://docs.sglang.ai/backend/structured_outputs.html)
Can you also add support to offline engine and native api?
Name python/sglang/srt/completion_template.py to code_completion_parser.py
Can you add support for streaming request?
Those parser like features are somehow similar, can you have a look of #2700 and #4000 ?
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T18:08:32Z
user: login: shuaills
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-03-13T21:30:42Z
user: login: woodx9
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-14T18:25:25Z
user: login: woodx9
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-14T18:25:30Z
user: login: shuaills
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-16T01:34:28Z

Review comments:
diff_hunk: @@ -0,0 +1,71 @@
+import unittest
+
+import openai
+
+from sglang.srt.hf_transformers_utils import get_tokenizer
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    popen_launch_server,
+)
+
+
+class TestOpenAIServer(unittest.TestCase):
path: test/srt/test_fim_completion.py
user: login: shuaills
body: TestFIMCompletion
created_at: 2025-03-03T17:55:38Z
author_association: COLLABORATOR
diff_hunk: @@ -172,6 +180,46 @@ def load_chat_template_for_openai_api(tokenizer_manager, chat_template_arg, mode
     # 2. For VLM, when chat_template_arg is None, set it automatically by guessing from model_path.
 
 
+def load_completion_template_for_openai_api(tokenizer_manager, completion_template_arg):
path: python/sglang/srt/openai_api/adapter.py
user: login: shuaills
body: Consider moving template loading and management logic from the adapter layer to the code_completion_parser.py module
The adapter should primarily handle API compatibility, not core functionality
created_at: 2025-03-13T21:14:26Z
author_association: COLLABORATOR
diff_hunk: @@ -504,7 +552,14 @@ def v1_generate_request(
                 "To compute logprobs of input prompt, please use the native /generate API."
             )
 
-        prompts.append(request.prompt)
+        if completion_template_name is None:
+            prompts.append(request.prompt)
+        else:
+            prompt = generate_completion_prompt_from_request(
+                request, completion_template_name
+            )
+            prompts.append(prompt)
path: python/sglang/srt/openai_api/adapter.py
user: login: shuaills
body: The request handler should not be concerned with template implementation details.
Move it to code_completion_parser.py
```
# in code_completion_parser.py
def prepare_prompt(request):
    template_name = get_current_template_name()
    
    if template_name is None or request.suffix is None:
        return request.prompt
    else:
        return generate_completion_prompt_from_request(request, template_name)
```

created_at: 2025-03-13T21:27:02Z
author_association: COLLABORATOR
diff_hunk: @@ -172,6 +180,46 @@ def load_chat_template_for_openai_api(tokenizer_manager, chat_template_arg, mode
     # 2. For VLM, when chat_template_arg is None, set it automatically by guessing from model_path.
 
 
+def load_completion_template_for_openai_api(tokenizer_manager, completion_template_arg):
path: python/sglang/srt/openai_api/adapter.py
user: login: woodx9
body: done
created_at: 2025-03-14T18:25:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -504,7 +552,14 @@ def v1_generate_request(
                 "To compute logprobs of input prompt, please use the native /generate API."
             )
 
-        prompts.append(request.prompt)
+        if completion_template_name is None:
+            prompts.append(request.prompt)
+        else:
+            prompt = generate_completion_prompt_from_request(
+                request, completion_template_name
+            )
+            prompts.append(prompt)
path: python/sglang/srt/openai_api/adapter.py
user: login: woodx9
body: done
created_at: 2025-03-14T18:25:30Z
author_association: CONTRIBUTOR

PR comments:
user: login: woodx9
created_at: 2025-02-26T05:49:03Z
body: @zhaochenyang20 cc
author_association: CONTRIBUTOR
user: login: zhaochenyang20
created_at: 2025-02-26T06:50:58Z
body: @woodx9 please rebase with main and add unit tests. following this:

https://docs.sglang.ai/references/contribution_guide.html

We will review this soon.
author_association: COLLABORATOR
user: login: shuaills
created_at: 2025-03-02T11:46:41Z
body: @woodx9 Can you fix the lint and add unit tests?
author_association: COLLABORATOR
user: login: woodx9
created_at: 2025-03-02T13:33:36Z
body: @shuaills Sure, will do
author_association: CONTRIBUTOR
user: login: woodx9
created_at: 2025-03-02T17:13:16Z
body: demo for this feature
`python -m sglang.launch_server --model-path deepseek-ai/deepseek-coder-1.3b-base --completion-template deepseek_coder --port 30020 --host 0.0.0.0`

author_association: CONTRIBUTOR
user: login: zhaochenyang20
created_at: 2025-03-02T17:33:16Z
body: @woodx9 fix the conflicts and @shuaills could you further review this?
author_association: COLLABORATOR
user: login: woodx9
created_at: 2025-03-03T05:09:29Z
body: will do tonight
author_association: CONTRIBUTOR
user: login: zhaochenyang20
created_at: 2025-03-03T08:29:50Z
body: @woodx9 Cool. Thnaks so much. You can reach out us to review on wechat also.
author_association: COLLABORATOR
user: login: woodx9
created_at: 2025-03-03T16:37:47Z
body: @zhaochenyang20 @shuaills I have rebased the master branch, plz review the code
author_association: CONTRIBUTOR
user: login: woodx9
created_at: 2025-03-03T17:25:23Z
body: @zhaochenyang20 sry, I forgot to run pre commit after rebase. plz run the test again.
author_association: CONTRIBUTOR
user: login: woodx9
created_at: 2025-03-03T18:19:18Z
body: I will do it. thank you @shuaills 
author_association: CONTRIBUTOR
user: login: zhaochenyang20
created_at: 2025-03-03T18:24:49Z
body: @woodx9 Have a good night~
author_association: COLLABORATOR
user: login: zhaochenyang20
created_at: 2025-03-07T19:48:21Z
body: @woodx9 I personally do not review this. Please ask shuai. And, please solve the confilcts.
author_association: COLLABORATOR
user: login: shuaills
created_at: 2025-03-08T14:18:26Z
body: Thanks! I'll review this today
author_association: COLLABORATOR
user: login: woodx9
created_at: 2025-03-08T14:20:36Z
body: hi, @shuaills. native api /generate don't support suffix parameter yet. Do we need to support this time? 
author_association: CONTRIBUTOR
user: login: woodx9
created_at: 2025-03-13T06:10:21Z
body: hi @shuaills  a kindly remind
author_association: CONTRIBUTOR
user: login: zhaochenyang20
created_at: 2025-03-13T08:49:57Z
body: i reminded him. thanks
author_association: COLLABORATOR
user: login: shuaills
created_at: 2025-03-13T15:34:33Z
body: Great work. Can we make docs a separated PR?
author_association: COLLABORATOR
user: login: woodx9
created_at: 2025-03-13T15:43:36Z
body: Sure, will do today~
author_association: CONTRIBUTOR
user: login: woodx9
created_at: 2025-03-13T16:26:03Z
body: Done for the separation of the code.
author_association: CONTRIBUTOR
user: login: shuaills
created_at: 2025-03-13T21:39:03Z
body: Feel free to reach out me on slack for discussion :)
author_association: COLLABORATOR
user: login: woodx9
created_at: 2025-03-14T16:22:49Z
body: will adapt it today
author_association: CONTRIBUTOR
user: login: zhaochenyang20
created_at: 2025-03-15T05:32:36Z
body: @woodx9 is it ready for review?
author_association: COLLABORATOR
user: login: woodx9
created_at: 2025-03-15T05:34:09Z
body: @zhaochenyang20 yes, @shuaills will review it tonight. thanks!
author_association: CONTRIBUTOR
user: login: zhaochenyang20
created_at: 2025-03-15T05:40:11Z
body: @woodx9 cool. ping me if needed
author_association: COLLABORATOR
user: login: zhaochenyang20
created_at: 2025-03-15T16:47:06Z
body: @shuaills if you agree, give an approval at https://github.com/sgl-project/sglang/pull/3612/files
author_association: COLLABORATOR
user: login: zhaochenyang20
created_at: 2025-03-16T02:23:55Z
body: @shuaills shuai, I think you have the access to rerun the CI. could you run it till finish on yourself if you approve it. no need to always rebase.
author_association: COLLABORATOR
user: login: woodx9
created_at: 2025-03-16T15:16:47Z
body: @zhaochenyang20 do we need to run CI again?
author_association: CONTRIBUTOR
user: login: zhaochenyang20
created_at: 2025-03-17T01:08:35Z
body: @woodx9 let me run it for you and merge it today.
author_association: COLLABORATOR