kvcache io kernels and test case

[xiezhq-hermann]

Motivation

This is a breakdown of #7313, upstreaming only changes on sgl-kernel and test cases.

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Summary of Changes

Hello @xiezhq-hermann, 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 introduces new, optimized kernel-level and direct I/O operations for Key-Value (KV) cache management within the sgl-kernel library. These operations facilitate efficient transfer of KV cache data across different memory locations or layouts, supporting both per-layer and all-layer transfers, as well as specialized multi-layer attention (MLA) formats. The changes include new CUDA kernels, their PyTorch bindings, Python wrappers, and comprehensive unit tests to ensure correctness and performance.

Highlights

• New KV Cache I/O Kernels: Implemented highly optimized CUDA kernels (transfer.cu) for efficient Key-Value cache data transfer. These kernels include transfer_item_warp for low-level data movement and transfer_kernel_impl for orchestrating transfers based on indices and memory layouts, supporting both non-MLA and multi-layer attention (MLA) formats.
• PyTorch Bindings and Python API: Exposed the new CUDA kernels to Python via PyTorch's C++ extension mechanism, registering functions like transfer_kv_per_layer, transfer_kv_all_layer, and their _mla and _direct counterparts. User-friendly Python functions (sgl_kernel/kvcacheio.py) abstract these implementations, allowing users to choose between 'kernel' (CUDA-optimized) and 'direct' (PyTorch copy_ based) I/O backends.
• Comprehensive Unit Tests: Added extensive unit tests (test_kvcacheio.py) using pytest to validate the correctness of the KV cache transfer operations. These tests cover various data types (bfloat16, float16), transfer sizes, page sizes, and configurations (MLA vs. non-MLA, per-layer vs. all-layer), comparing results against a reference Python implementation.

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

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 PR introduces KV cache I/O kernels, bindings, and tests. The kernels support transferring KV cache data with options for per-layer/all-layer and separate/combined K/V tensors, offering both a custom CUDA kernel and a direct PyTorch copy-based backend.

Feedback focuses on clarifying TODOs, improving naming, and addressing a potential performance bottleneck in the "direct" transfer path.

[Alisehen]

I have evaluated this PR using the HiCache benchmark, and it shows clear performance improvements.

For the multi-turn benchmark, I used the Qwen2.5-32B-INT8-TP2 model with 20 clients and 10 rounds.
• When the input length is 512, the PR achieves a 1.14× speedup in E2E latency and 1.4× in TTFT.
• When the input length is 1024, it achieves a 1.9× speedup in E2E and 2.5× in TTFT.

For the shared prefix benchmark, I also used Qwen2.5-32B-INT8-TP2, with 100 shuffled requests sampled from the Loogle dataset.
•With a shared prefix length of 4K, the PR improves E2E by 1.2× and TTFT by 1.4×.
•At 8K, the speedup is 1.8× for both E2E and TTFT.
•At 16K, the speedup is 1.2× for both E2E and TTFT.