[PD] Fix top logprobs crash in prefill path (#26299)
This commit is contained in:
@@ -510,6 +510,17 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
logits_output.input_token_logprobs = tuple(
|
||||
logits_output.input_token_logprobs.tolist()
|
||||
)
|
||||
if logits_output.next_token_top_logprobs_val:
|
||||
logits_output.next_token_top_logprobs_val = [
|
||||
v.tolist() for v in logits_output.next_token_top_logprobs_val
|
||||
]
|
||||
logits_output.next_token_top_logprobs_idx = [
|
||||
x.tolist() for x in logits_output.next_token_top_logprobs_idx
|
||||
]
|
||||
if logits_output.next_token_token_ids_logprobs_val:
|
||||
logits_output.next_token_token_ids_logprobs_val = [
|
||||
v.tolist() for v in logits_output.next_token_token_ids_logprobs_val
|
||||
]
|
||||
|
||||
for i, (req, next_token_id) in enumerate(
|
||||
zip(batch.reqs, next_token_ids, strict=True)
|
||||
|
||||
@@ -72,6 +72,32 @@ class TestDisaggregationAccuracy(PauseResumeInPlaceMixin, PDDisaggregationServer
|
||||
len(input_logprobs) > 0
|
||||
), f"input_logprobs should have at least one token, but got {len(input_logprobs)}"
|
||||
|
||||
def test_chat_completion_top_logprobs(self):
|
||||
client = openai.Client(api_key="empty", base_url=f"{self.lb_url}/v1")
|
||||
response = client.chat.completions.create(
|
||||
model="dummy",
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful AI assistant."},
|
||||
{"role": "user", "content": "What is the capital of France?"},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=8,
|
||||
logprobs=True,
|
||||
top_logprobs=5,
|
||||
)
|
||||
|
||||
self.assertIsNotNone(response.choices[0].logprobs)
|
||||
content_logprobs = response.choices[0].logprobs.content
|
||||
self.assertGreater(len(content_logprobs), 0)
|
||||
|
||||
first_top_logprobs = next(
|
||||
(item.top_logprobs for item in content_logprobs if item.top_logprobs),
|
||||
None,
|
||||
)
|
||||
self.assertIsNotNone(first_top_logprobs)
|
||||
self.assertGreater(len(first_top_logprobs), 0)
|
||||
self.assertIsInstance(first_top_logprobs[0].token, str)
|
||||
|
||||
def test_structured_output(self):
|
||||
json_schema = json.dumps(
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user