[Feature] Radix Tree in C++ #7369

DarkSharpness · 2025-06-20T00:02:15Z

Motivation

This PR is still part of #7194.

Modifications

Upstream Hierarchical Radix Tree in C++:
Added the core radix tree implementation for future hierarchical caching.
Python Interface (Non-Hierarchical Only):
Implemented Python bindings for the radix tree, currently supporting only non-hierarchical caching. Hierarchical cache integration will be addressed in subsequent PRs.
Experimental Feature Flag:
The experimental C++ radix tree can be enabled by setting the environment variable SGLANG_EXPERIMENTAL_HIRADIX_CACHE=1. Default behavior remains unchanged.

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 @DarkSharpness, 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 new, high-performance C++ Radix Tree implementation designed to serve as a core component for advanced memory caching, particularly for future hierarchical caching strategies. It includes the necessary Python bindings to expose this functionality, which can be activated on an experimental basis, and updates the build configuration to support the new C++ codebase.

Highlights

Core C++ Radix Tree Implementation: I've introduced a new C++-based Radix Tree (RadixTree) with foundational support for hierarchical caching concepts, including device (GPU) and host (CPU) memory management, eviction policies, and reference counting. This implementation lays the groundwork for future hierarchical caching features.
Python Bindings: I've provided Python bindings for the C++ Radix Tree (HiRadixCacheCpp) using pybind11. This allows Python code to interact with the high-performance C++ implementation. Currently, the Python interface primarily supports non-hierarchical caching, with full hierarchical integration planned for subsequent pull requests.
Experimental Feature Flag: The new C++ Radix Tree has been integrated into the scheduler as an experimental feature. Users can enable it by setting the environment variable SGLANG_EXPERIMENTAL_HIRADIX_CACHE=1. This allows for controlled rollout and testing.
Build System Integration: I've updated the CMake build system to compile the new C++ Radix Tree sources and generate the necessary Python extension module. This ensures seamless integration of the new C++ components into the project's build process.

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

This pull request introduces a C++ Radix Tree implementation with Python bindings, primarily for non-hierarchical caching. The C++ code is well-structured and uses modern features. The Python interface is clear.

The most critical missing piece is unit testing, which is essential for such a complex data structure. Documentation updates and benchmark results are also noted as pending in the PR description.

A few minor points regarding parameter naming consistency and assertions have been raised. The NotImplementedError for hierarchical cache support in the C++ backend is aligned with the PR's stated scope.

python/sglang/srt/managers/scheduler.py

python/sglang/srt/mem_cache/hiradix_cache_cpp.py

sgl-kernel/csrc/radix_tree/tree_v2_pool.h

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

KevinZeng08 · 2025-06-20T00:44:11Z

Great work!
I would like to ask how much the performance of the C++ version of RadixTree can be expected to improve?

DarkSharpness · 2025-06-20T05:00:00Z

Great work! I would like to ask how much the performance of the C++ version of RadixTree can be expected to improve?

Thank you for your question! Based on my testing, the C++ implementation typically delivers 2-3× speedups over the Python version. In extreme cases like #6907, this should significantly reduce evict operation overhead.

Note that:

This PR is part of [Refactor] Rewrite Hierarchical Prefix Cache in C++ #7194, which is not meant to replace the original radix cache in Python. Its ultimate goal is to accelerate some advanced prefix-matching related scheduling when hierarchical cache is enabled.
The C++ implementation is actually 10 times faster than Python Radix Cache when tested alone. But in production, some other operations like torch.cat will amortize the speedup.
Current Python implementation works fine for most workloads.

…ement)

merrymercy

move this to JIT https://docs.pytorch.org/tutorials/advanced/cpp_extension.html#jit-compiling-extensions. Only jit/import when it is really used. Open a new folder python/sglang/srt/mem_cache/cpp_radix_tree
Add an e2e test case test/srt/test_cpp_radix_tree.py. Follow this

sglang/test/srt/test_page_size.py

Line 16 in 8364608

class TestPageSize(CustomTestCase):
- Test e2e accuracy
- Test cache hit rate.

python/sglang/srt/managers/scheduler.py

DarkSharpness · 2025-07-05T07:44:41Z

All fixed. cc @merrymercy @xiezhq-hermann

DarkSharpness added 2 commits June 19, 2025 16:49

feat: upstream C++ hiradix tree implementation

994f13d

feat: support non-hierarchical part of hiradix cache C++ in Python

8cfddca

DarkSharpness requested review from zhyncs, ispobock, HandH1998, BBuf, yizhang2077, merrymercy, FlamingoPg, HaiShaw, Ying1123, hnyls2002 and xiezhq-hermann as code owners June 20, 2025 00:02

gemini-code-assist bot reviewed Jun 20, 2025

View reviewed changes

Update sgl-kernel/csrc/radix_tree/tree_v2_pool.h

55c8cb8

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

DarkSharpness added 2 commits June 19, 2025 18:36

style: clean up C++ code

3790408

minor: give better hint for assertion in C++

5371e7c

DarkSharpness and others added 7 commits June 21, 2025 15:30

fix: fix LRU eviction (since C++ priority queue top return largest el…

2bb5324

…ement)

minor: minor format and clean up code in C++

63df499

minor: make pre-commit happy

8ab6664

test: add a python test on C++ radix tree

1738ca8

Merge branch 'main' into dev_radix_cpp_impl

2a88daf

Merge branch 'main' into dev_radix_cpp_impl

a8110c7

Merge branch 'main' into dev_radix_cpp_impl

88d0c7f

merrymercy requested changes Jul 5, 2025

View reviewed changes

python/sglang/srt/managers/scheduler.py Outdated Show resolved Hide resolved

python/sglang/srt/managers/scheduler.py Outdated Show resolved Hide resolved

DarkSharpness added 2 commits July 4, 2025 22:34

chore: move C++ radix tree out of sgl_kernel and clean up

7a676e4

chore: copy .clang-format from sgl_kernel

b3a509b

DarkSharpness force-pushed the dev_radix_cpp_impl branch from 3a4a52f to b3a509b Compare July 5, 2025 05:38

DarkSharpness and others added 4 commits July 4, 2025 22:42

chore: revert all changes to in sgl-kernel

e6c2755

chore: add test and minor fixes

11bfa3f

refactor: remove most C++ hicache-related code

7d38ca6

Merge branch 'main' into dev_radix_cpp_impl

6f6436c

Merge branch 'main' into dev_radix_cpp_impl

bce98e1

xiezhq-hermann assigned xiezhq-hermann and merrymercy Jul 9, 2025

Merge branch 'main' into dev_radix_cpp_impl

1276cca

merrymercy mentioned this pull request Jul 21, 2025

Development Roadmap (2025 H2) #7736

Open

1 task

merrymercy added 3 commits August 2, 2025 18:46

Merge branch 'main' into dev_radix_cpp_impl

485e652

Use softlink for .clang-format

1714ae6

Add radix cache cpp

ddc2286

merrymercy approved these changes Aug 3, 2025

View reviewed changes

merrymercy merged commit e273aa6 into sgl-project:main Aug 3, 2025
57 of 62 checks passed

merrymercy mentioned this pull request Aug 3, 2025

Minor Optimizations in Radix Cache and Schedule Batch #6907

Closed

6 tasks

lifuhuang pushed a commit that referenced this pull request Aug 3, 2025

[Feature] Radix Tree in C++ (#7369)

f4b97c3

htiennv pushed a commit to htiennv/sglang that referenced this pull request Aug 5, 2025

[Feature] Radix Tree in C++ (sgl-project#7369)

522cdc6

ShangmingCai pushed a commit that referenced this pull request Aug 5, 2025

[Feature] Radix Tree in C++ (#7369)

2d8bf2c

ShangmingCai pushed a commit that referenced this pull request Aug 5, 2025

[Feature] Radix Tree in C++ (#7369)

36fc778

DarkSharpness deleted the dev_radix_cpp_impl branch August 10, 2025 22:39

narutolhy pushed a commit to narutolhy/sglang that referenced this pull request Aug 17, 2025

[Feature] Radix Tree in C++ (sgl-project#7369)

dda3d61

narutolhy pushed a commit to narutolhy/sglang that referenced this pull request Aug 18, 2025

[Feature] Radix Tree in C++ (sgl-project#7369)

91dbe84

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Feature] Radix Tree in C++ #7369

[Feature] Radix Tree in C++ #7369

Uh oh!

DarkSharpness commented Jun 20, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

KevinZeng08 commented Jun 20, 2025

Uh oh!

DarkSharpness commented Jun 20, 2025

Uh oh!

merrymercy left a comment •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

DarkSharpness commented Jul 5, 2025

Uh oh!

Uh oh!

Uh oh!

[Feature] Radix Tree in C++ #7369

[Feature] Radix Tree in C++ #7369

Uh oh!

Conversation

DarkSharpness commented Jun 20, 2025

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!

Uh oh!

Uh oh!

Uh oh!

KevinZeng08 commented Jun 20, 2025

Uh oh!

DarkSharpness commented Jun 20, 2025

Uh oh!

merrymercy left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

DarkSharpness commented Jul 5, 2025

Uh oh!

Uh oh!

Uh oh!

merrymercy left a comment •

edited

Loading