[BugFix] Fix cu_num_generated_tokens slicing logic in LogprobsLists.slice() method

[usberkeley]

Purpose

Problem Description

In LogprobsLists.slice(), directly slicing cu_num_generated_tokens using self.cu_num_generated_tokens[start_req_idx:end_req_idx] caused the sliced cumulative array to not start from 0, which could lead to indexing errors in subsequent operations.

Solution

Recalculate the sliced cu_num_generated_tokens to ensure it starts from 0:

1. Calculate the offset: cu_num_offset = self.cu_num_generated_tokens[start_req_idx]
2. Subtract the offset from each value in the sliced array to ensure the first element is 0

Test Plan

Add TestLogprobsLists(TestCase)

1. test_slice_without_cu_num_generated_tokens: Test slicing without cu_num_generated_tokens
2. test_slice_from_start: Test slicing from the start position
3. test_slice_from_middle: Test slicing from the middle position
4. test_slice_single_request: Test slicing a single request
5. test_slice_last_request: Test slicing the last request
6. test_slice_all_requests: Test slicing all requests (full slice)

Test Result

Pass

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request addresses a bug in the LogprobsLists.slice() method where cu_num_generated_tokens was not being correctly sliced and re-based, potentially leading to indexing errors. The proposed fix correctly recalculates the cumulative token counts for the sliced portion to ensure it starts from zero. The change is accompanied by a comprehensive new test suite in tests/v1/test_outputs.py which thoroughly validates the fix across various scenarios, including slicing from the start, middle, and end, as well as handling single and full slices. The logic of the fix is sound, and the tests are well-written, covering the relevant edge cases. Overall, this is a solid contribution that improves the correctness and robustness of the code.

[usberkeley]

Hi @hmellor
Could you help review this PR when you have time? Thanks!

[njhill]

Thanks @usberkeley!

[hmellor]

The new test file is not run in CI. @usberkeley please make a follow up PR ensuring that these tests are actually run.

[usberkeley]

The new test file is not run in CI. @usberkeley please make a follow up PR ensuring that these tests are actually run.

@hmellor Sorry about that. Sure, I’ll submit a new PR to fix this issue

[hmellor]

Don't worry, it's an easy mistake to make with the current CI setup. This should not have to be somthing contributors need to think about.

In future we hope to make this a non-issue.

[usberkeley]

Don't worry, it's an easy mistake to make with the current CI setup. This should not have to be somthing contributors need to think about.

In future we hope to make this a non-issue.

Thanks for the clarification! Glad to hear it’ll be easier in the future.