feat(api): add sglext_spec (#33518)

Signed-off-by: Muqi Li <muqi1029@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Muqi Li
2026-08-19 11:54:18 +08:00
committed by GitHub
co-authored by Codex
parent 3b065a56b0
commit 88f6074392
6 changed files with 360 additions and 10 deletions
@@ -43,6 +43,31 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
def _spec_result(index):
return {
"text": f"choice-{index}",
"meta_info": {
"id": "chatcmpl-spec-test",
"prompt_tokens": 10,
"completion_tokens": 2,
"cached_tokens": 0,
"finish_reason": {"type": "stop"},
"weight_version": "default",
"spec_accept_rate": 0.5,
"spec_accept_length": 2.0,
"spec_cap_length": index + 1.0,
"spec_block_accept_length": index + 0.5,
"spec_num_correct_drafts": 1,
"spec_num_proposed_drafts": 2,
"spec_verify_ct": 1,
"spec_correct_drafts_histogram": [0, 1],
"spec_cap_lens_histogram": [index, 1],
},
"index": index,
}
_DSV4_PREVIEW_ENCODER = 'REASONING_EFFORT_MAX = "preview"\n'
_DSV4_OFFICIAL_ENCODER = (
"REASONING_EFFORT_PROMPTS: Dict[str, str] = "
@@ -2494,6 +2519,34 @@ class ServingChatTestCase(unittest.TestCase):
},
)
def test_parallel_sampling_returns_spec_details_per_choice(self):
req = ChatCompletionRequest(
model="x",
messages=[{"role": "user", "content": "Hi?"}],
max_tokens=100,
n=2,
return_spec_tokens_details=True,
)
ret = [_spec_result(index) for index in range(2)]
response = self.chat._build_chat_response(req, ret, 1234567890)
details = response.sglext.spec_tokens_details
self.assertEqual([item.spec_cap_length for item in details], [1.0, 2.0])
self.assertEqual(
[item.spec_cap_lens_histogram for item in details],
[[0, 1], [1, 1]],
)
single_req = req.model_copy(update={"n": 1})
single_response = self.chat._build_chat_response(
single_req, ret[:1], 1234567890
)
self.assertEqual(
single_response.sglext.spec_tokens_details.spec_cap_length,
1.0,
)
def test_non_streaming_chat_response_returns_requested_token_ids_and_meta_info(
self,
):
@@ -2609,6 +2662,31 @@ class ServingChatTestCase(unittest.TestCase):
},
)
def test_streaming_parallel_sampling_orders_spec_details_by_choice(self):
async def mock_generate():
for index in (1, 0):
yield _spec_result(index)
self.tm.generate_request.return_value = mock_generate()
req = ChatCompletionRequest(
model="x",
messages=[{"role": "user", "content": "Hi?"}],
max_tokens=100,
n=2,
stream=True,
return_spec_tokens_details=True,
)
parsed = self._parse_chunks(self._run_chat_stream(Mock(), req))
details = next(chunk["sglext"] for chunk in parsed if "sglext" in chunk)[
"spec_tokens_details"
]
self.assertEqual([item["spec_cap_length"] for item in details], [1.0, 2.0])
self.assertEqual(
[item["spec_cap_lens_histogram"] for item in details],
[[0, 1], [1, 1]],
)
def _collect_continuous_usage(self, cached_tokens):
content = {
"text": "Hello",
@@ -25,6 +25,30 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
def _spec_result(index):
return {
"text": f"choice-{index}",
"meta_info": {
"id": "cmpl-spec-test",
"prompt_tokens": 10,
"completion_tokens": 2,
"cached_tokens": 0,
"finish_reason": {"type": "stop"},
"weight_version": "default",
"spec_accept_rate": 0.5,
"spec_accept_length": 2.0,
"spec_cap_length": index + 1.0,
"spec_block_accept_length": index + 0.5,
"spec_num_correct_drafts": 1,
"spec_num_proposed_drafts": 2,
"spec_verify_ct": 1,
"spec_correct_drafts_histogram": [0, 1],
"spec_cap_lens_histogram": [index, 1],
},
"index": index,
}
class _MockTemplateManager:
"""Minimal mock for TemplateManager."""
@@ -400,6 +424,103 @@ class ServingCompletionTestCase(unittest.TestCase):
},
)
def test_parallel_sampling_returns_spec_details_per_choice(self):
req = CompletionRequest(
model="x",
prompt="Hello world",
max_tokens=100,
n=2,
return_spec_tokens_details=True,
)
ret = [_spec_result(index) for index in range(2)]
response = self.sc._build_completion_response(req, ret, 1234567890)
details = response.sglext.spec_tokens_details
self.assertEqual(len(details), 2)
self.assertEqual(details[0].spec_cap_length, 1.0)
self.assertEqual(details[0].spec_block_accept_length, 0.5)
self.assertEqual(details[0].spec_cap_lens_histogram, [0, 1])
self.assertEqual(details[1].spec_cap_length, 2.0)
self.assertEqual(details[1].spec_block_accept_length, 1.5)
self.assertEqual(details[1].spec_cap_lens_histogram, [1, 1])
single_req = req.model_copy(update={"n": 1})
single_response = self.sc._build_completion_response(
single_req, ret[:1], 1234567890
)
self.assertEqual(
single_response.sglext.spec_tokens_details.spec_cap_length,
1.0,
)
disabled_req = single_req.model_copy(
update={"return_spec_tokens_details": False}
)
disabled_response = self.sc._build_completion_response(
disabled_req, ret[:1], 1234567890
)
self.assertIsNone(disabled_response.sglext)
def test_streaming_parallel_sampling_orders_spec_details_by_choice(self):
async def mock_generate(*args, **kwargs):
for index in (1, 0):
yield _spec_result(index)
self.sc.tokenizer_manager.generate_request = mock_generate
req = CompletionRequest(
model="x",
prompt="Hello world",
max_tokens=100,
n=2,
stream=True,
return_spec_tokens_details=True,
)
adapted_request, _ = self.sc._convert_to_internal_request(req)
async def run_stream(request):
return [
chunk
async for chunk in self.sc._generate_completion_stream(
adapted_request, request, self.fastapi_request
)
]
chunks = get_or_create_event_loop().run_until_complete(run_stream(req))
parsed = [
json.loads(chunk[len("data: ") :])
for chunk in chunks
if chunk.startswith("data: ") and chunk.strip() != "data: [DONE]"
]
details = next(chunk["sglext"] for chunk in parsed if "sglext" in chunk)[
"spec_tokens_details"
]
self.assertEqual([item["spec_cap_length"] for item in details], [1.0, 2.0])
self.assertEqual(
[item["spec_cap_lens_histogram"] for item in details],
[[0, 1], [1, 1]],
)
async def mock_single_generate(*args, **kwargs):
async for content in mock_generate():
if content["index"] == 0:
yield content
self.sc.tokenizer_manager.generate_request = mock_single_generate
single_req = req.model_copy(update={"n": 1})
single_chunks = get_or_create_event_loop().run_until_complete(
run_stream(single_req)
)
single_parsed = [
json.loads(chunk[len("data: ") :])
for chunk in single_chunks
if chunk.startswith("data: ") and chunk.strip() != "data: [DONE]"
]
single_details = next(
chunk["sglext"] for chunk in single_parsed if "sglext" in chunk
)["spec_tokens_details"]
self.assertIsInstance(single_details, dict)
def test_streaming_cached_tokens_details_emits_sglext(self):
"""Test that streaming completion responses emit cached token details in sglext."""