[Core][Bookkeeping] Update cu_num_accepted_tokens for all req_index

[Jialin]

Purpose

In this PR, we ensure cu_num_accepted_tokens is updated for all request_index. To avoid position skipped / shifted due to empty or None sampled_ids.

Test Plan & Test Result

> pytest tests/v1/sample/test_logprobs.py -k test_spec_decode_logprobs
> pytest tests/v1/sample/test_rejection_sampler.py

---

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 correctly addresses a bug in the bookkeeping logic for speculative decoding. By moving the update of cu_num_accepted_tokens to before the check for empty sampled_ids, it ensures that the cumulative token counts are accurate for all requests, including those that don't produce new tokens in a given step. This prevents potential indexing errors and position shifts. The implementation is clean and effectively resolves the described issue. I have no further suggestions.

[Jialin]

Include authors and reviewers of #26060 to confirm if it's the right update.

CC @TheEpicDolphin @22quinn @njhill

[TheEpicDolphin]

Thanks for catching this @Jialin, looks good to me!

[Jialin]

Gentle nudge @njhill @22quinn @yeqcharlotte @houseroad

[22quinn]

👍

[njhill]

Thanks for catching @Jialin!

Any chance we could extend test(s) added in #26060 to catch this case?

[Jialin]

Thanks for catching @Jialin!

Any chance we could extend test(s) added in #26060 to catch this case?

@njhill Skimmed throughput the tests added in #26060, but seems most of them are e2e testing, so it might be a bit hard to add unit test to cover this.

However, after this change, as we're only updating numpy arrays in bookkeeping now. A potential followup step is to introduce a numba jit function to further speed up bookkeeping. (And I could definitely add unit tests to cover this case after with the more self-contained jit function. WDYT?

[njhill]

Thanks for catching @Jialin!
Any chance we could extend test(s) added in #26060 to catch this case?

@njhill Skimmed throughput the tests added in #26060, but seems most of them are e2e testing, so it might be a bit hard to add unit test to cover this.

However, after this change, as we're only updating numpy arrays in bookkeeping now. A potential followup step is to introduce a numba jit function to further speed up bookkeeping. (And I could definitely add unit tests to cover this case after with the more self-contained jit function. WDYT?

@Jialin fair enough, I wasn't sure whether we could tweak the e2e test to trigger the bug. If that's nontrivial then sounds ok re leaving to future unit tests.