fix(scheduler): advance the input-logprob offset for skipped prefill requests

[EnesYilmazcode]

Motivation

I was reading the two prefill result loops side by side, because disaggregation/prefill.py:703 defines advance_logprob_pt() and calls it on all three of its skip paths, and the loop in batch_result_processor.py has no equivalent. It needs one.

The flat input-logprob array holds one block per request in the batch, whether or not that request asked for logprobs. _apply_prefill_logprobs is written that way on purpose: it advances logprob_pt outside its return_logprob guard, so the cursor tracks positions rather than interest. The producer agrees, building extend_logprob_pruned_lens_cpu from every sequence in the batch with the same extend_len - start_len formula the consumer uses.

The finished-in-a-mixed-batch and retracted branch at line 307 skips a request without advancing that cursor. Every later request in the batch then reads a shifted slice and is served another request's prompt logprobs.

Nothing raises. The check in logprob_result_processor.py validates the slice length, not its offset, and the shift is always backwards so it never runs off the end. input_top_logprobs_val[i] is indexed per request and stays correct while input_token_logprobs is sliced by the cursor and goes wrong, which is the diagnostic signature and probably why this has been quiet.

It bites when the skipped request itself asked for input logprobs. At the default logprob_start_len of -1 its block is zero rows and the skip is harmless. Overlap scheduling is on by default and retraction is ordinary under KV pressure, so no unusual flag is involved.

Reported in #36938.

Two open PRs touch this block without covering this path:

• #34153 applies the same consume-the-slice idea to the req.to_finish path. That is a pending finish, so req.finished() is false and it routes through the normal branch, not this skip.
• #36724 adds a deferred_mm_abort branch immediately above this continue and advances the cursor there, leaving this branch alone.

Modifications

Advance logprob_pt by the skipped request's block before the continue.

I used calculate_num_input_logprobs rather than subtracting the lengths directly because it keeps multi-item scoring batches correct, where only delimiter positions have logprobs. If you would rather keep the skip path cheap and do the plain subtraction, that is a one-line change and I am happy either way.

Added test/registered/unit/managers/test_batch_result_processor_logprobs.py, built on the harness in the existing test_batch_result_processor_hidden_states.py and registered as register_cpu_ci(est_time=2, suite="base-a-test-cpu"). It drives the real process_batch_result_prefill with a retracted request in position 0 and asserts the following request gets its own block.

Accuracy Tests

Two prefill requests, the first retracted, extend_input_len_per_req=[2, 3], extend_logprob_start_len_per_req=[0, 0], over a flat array of [-1.0, -2.0, -30.0, -31.0, -32.0]:

                  input_token_logprobs_val on the active request
before            [None,  -1.0,  -2.0]      the retracted request's values
after             [None, -30.0, -31.0]      its own

Checking out the previous revision of the source file while keeping the new test reproduces the first row.

test/registered/unit/managers/ on CPU: 572 passed before, 573 after, the extra one being this test. The 2 failures and 3 errors in that directory (test_mm_process_config.py, test_customized_info_streaming.py) are present either way.

Speed Tests and Profiling

N/A. One integer add per skipped request, and only when batch.return_logprob is set.

Checklist

• Format the code with pre-commit. isort 7.0.0, black 26.1.0 and ruff 0.15.1 are clean on both files.
• Added a unit test that fails without the fix.
• The change is scoped to the reported defect.

The new test runs on CPU in about 8 seconds and needs no GPU.

---

CI States

Latest PR Test (Base): ❌ Run #36201305886
Latest PR Test (Extra): ❌ Run #36201305577
Latest PR Test (AMD ROCm 10): ❌ Run #36201305890

[rwang5203]

On current main, this fixture lacks sampling_mask_output, so the regression test fails before checking logprob offsets. Consider adding the current optional output fields to the fixture.

[EnesYilmazcode]

Thanks, fixed in 488b56a. Rebased on main and switched the fixture to the real GenerationBatchResult and LogitsProcessorOutput dataclasses, so new optional fields pick up their defaults instead of breaking the test. It passes on current main and still fails with the scheduler fix reverted.

[EnesYilmazcode]

@rwang5203 could you add the run-ci label when you get a chance? The gate is blocking the test run.