support memory pool page first layout by huangtingwei9988 · Pull Request #8061 · sgl-project/sglang

huangtingwei9988 · 2025-07-15T13:32:27Z

Motivation

Currently, sglang's memory_pool and memory_pool_host are both layer_first layouts, which is conducive to layer wise operations, for example, loading KV caches from host memory to device memory layer by layer.

    def init_kv_buffer(self):
        return torch.empty(
            (
                self.layer_num,
                self.size,
                1,
                self.kv_lora_rank + self.qk_rope_head_dim,
            ),
            dtype=self.dtype,
            device=self.device,
            pin_memory=self.pin_memory,
        )

However, when integrating with a remote kv cache pool like mooncake store, the layer_first layout will make it impossible to obtain the complete kv cache in the continuous address space. （For example, the layer0 caches of all tokens are together, but the layer0 caches of each token are not together with their layer1 caches.）

Therefore, when zero copying, it can only obtain the kv cache by layer. This will cause the number of keys in mooncake store to become layer_num times, which will reduce performance and make it difficult to manage for mooncake store

Co-author @AniZpZ @zhaoyongke @zhangzuo21

Modifications

support memory pool page first layout

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 @huangtingwei9988, 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 support for a 'page-first' memory layout for KV caches, addressing performance and management challenges encountered when integrating with remote KV cache systems that benefit from contiguous memory blocks. It updates the memory pool management and KV cache transfer mechanisms to be layout-aware, enabling efficient data movement between different memory organizations at the kernel level.

Highlights

New Memory Layout Support: Introduced a 'page-first' memory layout for KV caches, alongside the existing 'layer-first' layout, to improve integration with remote KV cache pools like Mooncake store.
Flexible KV Cache Transfer: Enhanced KV cache transfer functions (both per-layer and all-layer, and MLA/non-MLA variants) to handle transfers between different memory layouts (layer-first to layer-first, and page-first to layer-first).
Layout Dimension Tracking: Added layout_dim properties to memory_pool and memory_pool_host classes to accurately track the dimension of a single layer/page, which is crucial for calculating offsets in the new layout.
Kernel-level Layout Awareness: Modified underlying CUDA kernels to accept parameters indicating source and destination memory layouts, allowing them to dynamically select the correct global offset calculation logic (get_global_offset_lf for layer-first, get_global_offset_pf for page-first).
Expanded Testing: Added new test cases to verify the correctness of KV cache transfers, specifically including scenarios for converting from a page-first layout to a layer-first layout.

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.

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

The pull request introduces support for a page-first memory layout, which improves integration with remote KV cache pools. The changes include modifications to Python, C++, and CUDA code to handle different memory layouts. The new functionality is well-tested. I have provided some suggestions to improve readability and maintainability.

xiezhq-hermann · 2025-08-03T03:10:02Z

replaced by #8651

huangtingwei9988 added 11 commits July 14, 2025 19:06

support page_first kernel

96e9cce

fix

181794f

fix

26e0dc0

add unit test

c9e19d4

add page first to layer first copy unit test

4843893

fix

1ba2e21

fix

bb381c7

fix

dbb0c81

fix template args

4e1a1e0

fix

8a0ad90

fix

56ecb4e

huangtingwei9988 requested review from BBuf, FlamingoPg, HaiShaw, HandH1998, Ying1123, hnyls2002, ispobock, merrymercy, xiezhq-hermann, yizhang2077 and zhyncs as code owners July 15, 2025 13:32

gemini-code-assist Bot reviewed Jul 15, 2025

View reviewed changes

Comment thread sgl-kernel/csrc/kvcacheio/transfer.cu

Comment thread sgl-kernel/csrc/kvcacheio/transfer.cu

refactor

f861882

AniZpZ force-pushed the host_memory_page_first branch from 8f31c9e to f861882 Compare July 15, 2025 16:27

fmt

4a4ad2d

xiezhq-hermann self-assigned this Jul 15, 2025

Merge branch 'main' into host_memory_page_first

732161a

xiezhq-hermann mentioned this pull request Jul 24, 2025

Interface change for kvcache io to support page first layout #8318

Merged

6 tasks

huangtingwei9988 mentioned this pull request Aug 1, 2025

Support page first layout zero copy for mooncake store #8651

Merged

7 tasks

xiezhq-hermann closed this Aug 3, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

support memory pool page first layout#8061

support memory pool page first layout#8061
huangtingwei9988 wants to merge 14 commits intosgl-project:mainfrom
AniZpZ:host_memory_page_first

huangtingwei9988 commented Jul 15, 2025 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

xiezhq-hermann commented Aug 3, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

huangtingwei9988 commented Jul 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

xiezhq-hermann commented Aug 3, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

huangtingwei9988 commented Jul 15, 2025 •

edited

Loading