Skip to content

[Perf] openai: Tokenize dsv4 / dsv41 / dsv32 prompts with parallel_prompt_encode #45410

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kano-zheng wants to merge 1 commit into from
Closed

[Perf] openai: Tokenize dsv4 / dsv41 / dsv32 prompts with parallel_prompt_encode #45410

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions python/sglang/srt/entrypoints/openai/serving_chat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1667,7 +1667,9 @@ def _apply_jinja_template(
reasoning_effort=v4_reasoning_effort,
reasoning_effort_profile=reasoning_effort_profile,
)
prompt_ids = self.tokenizer_manager.tokenizer.encode(real_input)
prompt_ids = parallel_prompt_encode(
self.tokenizer_manager.tokenizer, real_input
)
elif is_dsv41:
if request.task is not None:
encoding_dsv41.attach_task_to_last_user_message(
Expand All @@ -1694,12 +1696,16 @@ def _apply_jinja_template(
self.tokenizer_manager.image_token_id
),
)
prompt_ids = self.tokenizer_manager.tokenizer.encode(real_input)
prompt_ids = parallel_prompt_encode(
self.tokenizer_manager.tokenizer, real_input
)
else:
real_input = encoding_dsv32.encode_messages(
messages, thinking_mode=thinking_mode
)
prompt_ids = self.tokenizer_manager.tokenizer.encode(real_input)
prompt_ids = parallel_prompt_encode(
self.tokenizer_manager.tokenizer, real_input
)

# Append assistant prefix if continue_final_message is enabled
if assistant_prefix:
Expand Down
28 changes: 28 additions & 0 deletions test/registered/unit/entrypoints/openai/test_serving_chat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1572,6 +1572,34 @@ def test_dsv_tool_choice_none_omits_tools_from_prompt(self):
prompt = self.tm.tokenizer.encode.call_args.args[0]
self.assertEqual("calculator" in prompt, expects_tools)

def test_native_encoders_tokenize_through_parallel_prompt_encode(self):
# dsv4 / dsv41 / dsv32 render the whole conversation to one string and
# then tokenize it on the HTTP event loop. Agent prompts reach several
# hundred thousand tokens, so these paths must use the same chunked
# encoder as the Jinja path instead of a single tokenizer.encode call.
self.template_manager.chat_template_name = None
self.template_manager.jinja_template_content_format = "string"
self.chat._dsv4_reasoning_effort_profile = "preview"

for chat_encoding_spec in ("dsv4", "dsv32", "dsv41"):
with self.subTest(chat_encoding_spec=chat_encoding_spec):
self.chat.chat_encoding_spec = chat_encoding_spec
request = ChatCompletionRequest(
model="x",
messages=[{"role": "user", "content": "Solve 100 / 7."}],
)
with patch(
"sglang.srt.entrypoints.openai.serving_chat.parallel_prompt_encode",
return_value=[7, 8, 9],
) as encode_mock:
result = self.chat._process_messages(request, is_multimodal=False)

encode_mock.assert_called_once()
tokenizer, prompt = encode_mock.call_args.args
self.assertIs(tokenizer, self.tm.tokenizer)
self.assertIn("Solve 100 / 7.", prompt)
self.assertEqual(result.prompt_ids, [7, 8, 9])

def test_kimi_k3_encoder_receives_wire_request_fields(self):
self.template_manager.chat_template_name = None
self.chat.chat_encoding_spec = "kimi_k3"
Expand Down
Loading