Use dynamically maintained num_waiting_tokens in get_load()

[vipwangerxiao]

Motivation

Use dynamically maintained num_waiting_tokens in get_load() to avoid list access

Modifications

• for no PD disaggregation:

num_waiting_tokens represents the total number of tokens in the waiting queue

It increases when requests enter the waiting queue, and decreases when requests leave the waiting queue to running list or abort

• for PD disaggregation P node:

num_waiting_tokens represents the total number of tokens in the bootstrap/waiting queue

It increases when requests enter the bootstrap queue, and decreases when requests leave the waiting queue to running list or fail in bootstrap queue

• for PD disaggregation D node:

num_waiting_tokens represents the total number of tokens in the prealloc/retract queue, not including waiting queue because KV cache allocating is done before request gets in waiting queue

It increases when requests enter the prealloc/retrac queue, and decreases when requests leave the prealloc queue to transfer queue, or leave retract queue to waiting queue, or abort in prealloc queue

Accuracy Tests

Compare the dynamically maintained num_waiting_tokens(num_tokens) with the value obtained directly from the queues(direct_sum_num_tokens), and they are expected to be the same.

no PD disaggregation

[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":261865,"num_reqs":87,"num_waiting_reqs":82,"num_tokens":261865}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":584829,"num_reqs":192,"num_waiting_reqs":190,"num_tokens":584829}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":585199,"num_reqs":195,"num_waiting_reqs":190,"num_tokens":585199}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":569844,"num_reqs":190,"num_waiting_reqs":185,"num_tokens":569844}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":554854,"num_reqs":182,"num_waiting_reqs":180,"num_tokens":554854}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":555204,"num_reqs":185,"num_waiting_reqs":180,"num_tokens":555204}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":539505,"num_reqs":180,"num_waiting_reqs":175,"num_tokens":539505}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":524590,"num_reqs":172,"num_waiting_reqs":170,"num_tokens":524590}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":524955,"num_reqs":175,"num_waiting_reqs":170,"num_tokens":524955}]

PD disaggregation: P node

[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":519089,"num_reqs":169,"num_waiting_reqs":169,"num_tokens":519089}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":560854,"num_reqs":184,"num_waiting_reqs":184,"num_tokens":560854}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":539500,"num_reqs":177,"num_waiting_reqs":177,"num_tokens":539500}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":521602,"num_reqs":170,"num_waiting_reqs":170,"num_tokens":521602}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":500625,"num_reqs":164,"num_waiting_reqs":164,"num_tokens":500625}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":484913,"num_reqs":159,"num_waiting_reqs":159,"num_tokens":484913}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":463922,"num_reqs":152,"num_waiting_reqs":152,"num_tokens":463922}]

PD disaggregation: D node

[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":6000,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":6000}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":0,"num_reqs":0,"num_waiting_reqs":0,"num_tokens":0}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":906911,"num_reqs":2,"num_waiting_reqs":0,"num_tokens":906911}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":1035255,"num_reqs":43,"num_waiting_reqs":41,"num_tokens":1035255}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":1014128,"num_reqs":39,"num_waiting_reqs":36,"num_tokens":1014128}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":998796,"num_reqs":35,"num_waiting_reqs":31,"num_tokens":998796}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":983560,"num_reqs":30,"num_waiting_reqs":26,"num_tokens":983560}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":968714,"num_reqs":25,"num_waiting_reqs":21,"num_tokens":968714}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":950668,"num_reqs":19,"num_waiting_reqs":15,"num_tokens":950668}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":930902,"num_reqs":14,"num_waiting_reqs":10,"num_tokens":930902}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":918929,"num_reqs":7,"num_waiting_reqs":4,"num_tokens":918929}]
[{"rid":null,"http_worker_ipc":null,"direct_sum_num_tokens":901005,"num_reqs":3,"num_waiting_reqs":0,"num_tokens":901005}]

Benchmarking and Profiling

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.

[gemini-code-assist]

Summary of Changes

Hello @vipwangerxiao, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refactors the token load calculation within the scheduler by introducing a dedicated counter for waiting tokens. This counter is updated incrementally as requests move through different queues, eliminating the need for expensive on-the-fly summations. The primary goal is to enhance performance and reduce computational overhead when determining the system's current load.

Highlights

• Performance Optimization: The pull request introduces a dynamically maintained num_waiting_tokens counter to replace the previous method of calculating total waiting tokens by iterating through various request queues. This change aims to improve the efficiency of the get_load() function by avoiding repeated list access.
• Token Tracking Logic: The num_waiting_tokens variable now accurately reflects the total number of tokens across the waiting, bootstrap, and prealloc/transfer queues. It is incremented when requests enter these queues and decremented when requests leave them (e.g., move to a running state, are preempted, or fail).
• Accuracy Verification: The changes have been verified with accuracy tests, demonstrating that the dynamically maintained num_waiting_tokens consistently matches the value obtained by directly summing tokens from the queues, ensuring correctness.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a significant performance improvement by replacing the expensive, on-demand calculation of num_waiting_tokens with a dynamically maintained counter. The implementation is well-executed, with the counter being correctly updated across various request lifecycle events and queues. This change not only optimizes the get_load() function but also fixes a pre-existing bug where tokens in the transfer queue were not being counted. I've also identified and suggested a fix for a similar issue in the calculation of num_waiting_reqs to ensure load balancing accuracy. Overall, this is a solid contribution that enhances both performance and correctness.