Add GDS alternative for hierarchical kv cache

[didoteebin]

Motivation

Add GDS alternative for hierarchical kv cache

Add options for GPU kvcache direct store to SSD , potential TTFT benefit may gain.

Modifications

python/sglang/srt/mem_cache/memory_pool_host.py to add GDS support in MHATokenToKVPoolHost
python/sglang/srt/mem_cache/gds.py to add a GDS class based on NIXL lib
python/sglang/global_config.py to add a switch for user to decide whether to enable gds or not

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 @didoteebin, 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 a significant enhancement to the hierarchical KV cache management by integrating GPU Direct Storage (GDS) capabilities. The primary goal is to optimize the "Time To First Token" (TTFT) by enabling direct data transfers between GPU memory and SSDs for KV cache pages, thereby reducing reliance on traditional host memory for these operations. This change provides a configurable option for users to leverage GDS, particularly within the multi-head attention (MHA) token pool, to potentially improve performance for large language models.

Highlights

• GPU Direct Storage (GDS) Integration: Implemented support for GPU Direct Storage (GDS) to enable direct KV cache transfers between GPU memory and SSDs, aiming to reduce Time To First Token (TTFT) by offloading data from host memory.
• GDS Configuration: Introduced enable_gds and gds_path options in global_config.py, allowing users to enable and configure GDS functionality via environment variables.
• Hierarchical KV Cache Enhancements: Modified MHATokenToKVPoolHost to conditionally leverage GDS for writing and loading KV cache pages across layers, providing an alternative to traditional host memory transfers.
• Memory Pool Abstraction: Added new abstract methods (transfer, get_flat_data, get_flat_data_by_layer, assign_flat_data) to the HostKVCache base class, formalizing the interface for memory operations within the hierarchical KV cache system.

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 pull request introduces a GPU Direct Storage (GDS) alternative for the hierarchical KV cache. My review identified a critical issue with incorrect function arguments, a high-severity issue with improper handling of an environment variable, and another high-severity issue regarding a potential resource leak. Additionally, there are several medium-severity concerns related to code clarity and feature completeness.

[gemini-code-assist]

The logger.info calls added in this __init__ method (e.g., lines 60, 63, 75, 91) appear to be for debugging purposes. They are either uninformative (e.g., "xxxx") or contain typos and extra characters (e.g., line 91 has "devcie" and "##...###"). Remove or change these logs to logger.debug to avoid cluttering production logs and improve code clarity.

[gemini-code-assist]

The GDS logic has been added for MHATokenToKVPoolHost, but not for MLATokenToKVPoolHost. The new methods in MLATokenToKVPoolHost only implement the non-GDS path. This creates an inconsistency where GDS will not be used for MLA-style models even if enabled. Add the GDS implementation for MLATokenToKVPoolHost to ensure consistent behavior.

[msharmavikram]

@didoteebin can you enable gds via dynamo/nixl? This would enable not just gds but all backend.

This should help significantly compared to just gds that allows file system.

[xinranwang17]

Great job! Is it possible to share some benchmark data here? Also, I wonder is there any plan to use disk GDS access as L3 cache instead of replacing L2 cpu memory directly?

[didoteebin]

@xinranwang17 Thanks , working on the benchmark to show TTFT benefit.
@xiezhq-hermann Question on the loading process .It seems that current logic is load whatever kvcache in cpu specified by host_indices to GPU area which specified by device_indices, there is no visibility of these kv cache will be eventually match or not when doing this loading process , am I right ?

[xiezhq-hermann]

@xinranwang17 Thanks , working on the benchmark to show TTFT benefit. @xiezhq-hermann Question on the loading process .It seems that current logic is load whatever kvcache in cpu specified by host_indices to GPU area which specified by device_indices, there is no visibility of these kv cache will be eventually match or not when doing this loading process , am I right ?

@didoteebin the loading would only start after matching the tokens if that's your question

[sohil-bst]

Hello, is there an update here? I'm looking for a solution to use my local ssd to extend kv cache. @didoteebin

[msharmavikram]

You can use SGL hicache with NIXL to get GDS support.

[sohil-bst]

I see both POSIX and GDS available for NIXL and plan to try both:

python3 -c "from nixl._api import nixl_agent, nixl_agent_config; \
    agent = nixl_agent('test', nixl_agent_config(backends=[])); \
    print(agent.get_plugin_list())"

_api.py:251 Initialized NIXL agent: test
['GDS', 'GDS_MT', 'GPUNETIO', 'GUSLI', 'OBJ', 'POSIX', 'UCX']

This PR led me to believe GDS was still a WIP though.

[xiezhq-hermann]

close for no update, but feel free to re-open it