[Kernel] refactor cpu worker v0 cache dtype #20080
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
There was a problem hiding this comment.
Summary of Changes
Hello @andyxning, I'm Gemini Code Assist1! 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 focuses on refactoring the CPU worker's KV cache data type determination logic to improve consistency and reduce code duplication. By centralizing the dtype resolution into a dedicated static method, I've ensured that both cache allocation and block size calculations use the same, unified logic, particularly for fp8 and fp8_e5m2 types.
Highlights
- Refactor KV Cache Dtype Logic: I've extracted the logic for determining the KV cache data type into a new static method,
CPUCacheEngine.get_kv_cache_dtype. This method centralizes the decision-making forauto,fp8, andfp8_e5m2cache types. - Unify Dtype Determination: The
CPUCacheEngine.__init__method now utilizes the newly introducedget_kv_cache_dtypeto setself.dtype, ensuring consistency with how the cache block size is calculated. - Consistent Cache Block Size Calculation: The
CPUCacheEngine.get_cache_block_sizemethod has been updated to accept theCacheConfigobject directly and now uses the sharedget_kv_cache_dtypemethod to determine thedtypefor calculating the block size in bytes. This eliminates redundant dtype logic. - Simplified Parameter Passing: The
CPUWorker.get_cache_block_size_bytesmethod now passes the entireself.cache_configobject toCPUCacheEngine.get_cache_block_size, aligning with the refactored signature of that method.
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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| 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 issue 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 is currently in preview and 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 to provide feedback.
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
-
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. ↩
There was a problem hiding this comment.
Code Review
The pull request effectively refactors the KV cache dtype determination logic in cpu_worker.py by centralizing it into a new static method. This improves code consistency and maintainability by eliminating duplication and simplifying function signatures. The changes are well-implemented and enhance the overall code quality.
Signed-off-by: Andy Xie <andy.xning@gmail.com>
2b6be2c to
211dcdb
Compare
|
/cc @bigPYJ1151 |
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Signed-off-by: Andy Xie <andy.xning@gmail.com> Signed-off-by: avigny <47987522+avigny@users.noreply.github.com>
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Signed-off-by: Andy Xie <andy.xning@gmail.com> Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com>
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Purpose
For now, the cpu worker v0 cache dtype gen logic are existed in both
__init__andget_cache_block_size. One is used in_allocate_kv_cacheand the other is used as the block size in bytes to compute the available cpu blocks. These two places should use the same logic to generate the cache dtype and keep consistent.For now, it is lucky that
torch.float8_e5m2andtorch.uint8are the same in size, they both are 1 bytes. But as the above explanation, the cache dtype should be a consistent one forfp8andfp8_e5m2Test Plan
NA
Test Result
NA
(Optional) Documentation Update