Reviewing files that changed from the base of the PR and between 7bc70bc and 222844d.

• lib/llm/src/kv_router.rs (3 hunks)
• lib/llm/src/kv_router/scheduler.rs (1 hunks)
• lib/llm/src/kv_router/sequence.rs (7 hunks)

• GitHub Check: Mirror Repository to GitLab
• GitHub Check: pre-merge-rust (.)
• GitHub Check: pre-merge-rust (lib/runtime/examples)
• GitHub Check: pre-merge-rust (lib/bindings/python)
• GitHub Check: Build and Test - vllm

lib/llm/src/kv_router/scheduler.rs (1)

262-266: LGTM: Clean batch token push implementation

The method signature change from single token to token slice is consistent with the PR objective. The implementation correctly forwards the slice to the underlying sequences layer.

lib/llm/src/kv_router.rs (2)

209-212: LGTM: Consistent API change for batch token push

The method signature change aligns with the batch token push implementation across the codebase.

---

276-320: Well-implemented batch token accumulation logic

The refactoring correctly implements batch token pushing by:

• Tracking initial sequence length and block boundaries
• Accumulating tokens from response chunks
• Pushing accumulated tokens only when crossing block boundaries
• Properly updating block indices after each push

The block boundary detection logic is sound - using (isl + total_output_length).saturating_sub(1) / block_size to calculate the current block index and comparing with the previous index.

Note: Any remaining accumulated tokens at stream end are handled by the free() call, which should properly clean up partial blocks.

lib/llm/src/kv_router/sequence.rs (5)

188-234: Excellent refactoring to handle batch token pushes with proper borrow management

The implementation correctly handles the transition from single-token to batch-token pushes:

1. Proper borrow management: The deferred operations pattern elegantly solves borrow checker issues by collecting operations during iteration and executing them after releasing the token_seq borrow.

2. Correct block boundary detection: Using token_seq.total_tokens() % self.block_size != 1 correctly identifies when a block boundary is crossed (first token in a new block).

3. Efficient batch processing: The method processes all tokens in the slice while maintaining the same block management semantics as the original single-token version.

The logic flow is clear and maintains data integrity throughout the batch operation.

---

248-249: Consistent enum variant update for batch token support

The UpdateSequences::Push variant correctly changed from single token to Vec<u32> to support the batch token push functionality.

---

311-313: Proper worker thread handling for batch tokens

The worker thread correctly processes the updated enum variant by passing the tokens slice to the underlying push method.

---

414-426: Well-implemented multi-worker batch token push

The method correctly:

• Accepts token slice parameter
• Finds the appropriate worker for the request
• Converts slice to Vec<u32> for sending across thread boundaries
• Maintains proper error handling

The slice-to-vector conversion is necessary for thread communication.

---

519-519: Good test updates for batch token push API

The test cases correctly demonstrate the new batch push functionality by pushing multiple tokens at once instead of individual token pushes.

Also applies to: 573-573