Skip to content

[Refactor] remove torch dependency from cache kernel - #2086

Draft
DwyaneShi wants to merge 1 commit into
vllm-project:mainfrom
DwyaneShi:haiyang/rm-torch-dep-kernel
Draft

[Refactor] remove torch dependency from cache kernel#2086
DwyaneShi wants to merge 1 commit into
vllm-project:mainfrom
DwyaneShi:haiyang/rm-torch-dep-kernel

Conversation

@DwyaneShi

Copy link
Copy Markdown
Collaborator

Pull Request Description

  • remove torch dependency from cache kernel
  • remove vllm dependency from cache ops unit testing

Related Issues

Resolves: #[Insert issue number(s)]

Important: Before submitting, please complete the description above and review the checklist below.


Contribution Guidelines (Expand for Details)

We appreciate your contribution to aibrix! To ensure a smooth review process and maintain high code quality, please adhere to the following guidelines:

Pull Request Title Format

Your PR title should start with one of these prefixes to indicate the nature of the change:

  • [Bug]: Corrections to existing functionality
  • [CI]: Changes to build process or CI pipeline
  • [Docs]: Updates or additions to documentation
  • [API]: Modifications to aibrix's API or interface
  • [CLI]: Changes or additions to the Command Line Interface
  • [Misc]: For changes not covered above (use sparingly)

Note: For changes spanning multiple categories, use multiple prefixes in order of importance.

Submission Checklist

  • PR title includes appropriate prefix(es)
  • Changes are clearly explained in the PR description
  • New and existing tests pass successfully
  • Code adheres to project style and best practices
  • Documentation updated to reflect changes (if applicable)
  • Thorough testing completed, no regressions introduced

By submitting this PR, you confirm that you've read these guidelines and your changes align with the project's contribution standards.

@DwyaneShi
DwyaneShi requested a review from Jeffwan April 7, 2026 22:44

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request refactors the KV cache operations to remove the dependency on torch.ops and vllm by transitioning to a ctypes-based interface for CUDA kernels. Key changes include updating the build system to use CUDAToolkit directly, rewriting CUDA kernels to use raw pointers and a C-compatible interface, and implementing local test utilities. Review feedback identified a critical bug in the data type dispatch logic for FP8 offloading, an incorrect variable name in the CMake configuration, and potential issues with platform portability and silent precision loss for FP64 types.

Comment thread python/aibrix_kvcache/csrc/cache_kernels.cu
Comment thread python/aibrix_kvcache/CMakeLists.txt Outdated
Comment thread python/aibrix_kvcache/aibrix_kvcache/_custom_ops.py Outdated
Comment thread python/aibrix_kvcache/aibrix_kvcache/_custom_ops.py
@DwyaneShi
DwyaneShi marked this pull request as draft April 8, 2026 00:11
@DwyaneShi
DwyaneShi force-pushed the haiyang/rm-torch-dep-kernel branch 4 times, most recently from 491d4fd to 4553710 Compare April 10, 2026 04:50
Signed-off-by: Haiyang Shi <haiyang.shi@bytedance.com>
@DwyaneShi
DwyaneShi force-pushed the haiyang/rm-torch-dep-kernel branch from 4553710 to ff54357 Compare April 10, 2026 05:04
@Jeffwan

Jeffwan commented Apr 12, 2026

Copy link
Copy Markdown
Collaborator

@DwyaneShi is this still WIP

@DwyaneShi

Copy link
Copy Markdown
Collaborator Author

@DwyaneShi is this still WIP

Yes, haven't conducted w/ e2e testing yet

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants