[scheduler] fix: correcting extend_logprob_start_len calculation

[ch-wan]

Motivation

The original calculation for extend_logprob_start_len was incorrect (code). When kv retraction happens, it should be determined based on len(fill_ids) (the total number of original input ids and output ids) when return_logprob is False but logprob_start_len is set as len(origin_input_ids) - 1.

Modifications

We set logprob_start_len as -1 by default when return_logprob is False so that extend_logprob_start_len can be set as len(fill_ids) - 1 during prepare_for_extend. We also reverted changes in #12115 as it's no longer needed.

Accuracy Tests, Benchmarking and Profiling

Tests in #9748 and #12115 are reproduced for verifying the robustness, accuracy and performance of this PR.

Test test_logprobs.py:

max Δ=0
mean Δ=0
logprobs returned for 500 samples (expected: 500)

✅ Test completed successfully!

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[ch-wan]

/tag-and-rerun-ci

[ch-wan]

it should be removed before this PR getting merged

[ch-wan]

/rerun-failed-ci

[merrymercy]

add a test case so the old code will fail / crash

[sundar24295s]

In the previous code, we didn't propogate logprob_start_len = -1 beyond the Scheduler. The value was reset to len(req.origin_input_ids) - 1. Hence, for prefill-only, we followed the same way.

Given, with this change we are propogating logprob_start_len = -1 into scheduler_batch, we can do the same logic for prefill_only as well instead of resetting to len(req.origin_input_ids) in line: 1531.

But, looks like, the current change doesn't break the prefill-only logprobs computation, so we are good.