[HiCache] Support host memory size smaller than device memory size by stmatengss · Pull Request #16909 · sgl-project/sglang

stmatengss · 2026-01-11T16:54:13Z

Motivation

To support host memory size can be smaller than device memory size:

In the evict method, adding a check for the return value of write_backup. If write_backup fails (e.g., due to host OOM), the code now gracefully falls back to _evict_regular, which evicts the node from the device without attempting to back it up.
Modifying evict_host to allow evicting host memory even for nodes that are still resident on the device (x.evicted is False). This is crucial when the host cache is smaller than the device cache, as we need to reclaim host memory occupied by redundant copies to make space for new backups.

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review Process

Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
After green CI and required approvals, ask Merge Oncalls to merge.

gemini-code-assist · 2026-01-11T16:54:31Z

Summary of Changes

Hello @stmatengss, 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 enhances the HiCache system's flexibility by allowing the host memory size to be smaller than the device memory size. This is achieved through refined eviction policies that handle host memory limitations more robustly, including graceful fallback for backup failures and enabling host memory reclamation for device-resident nodes. Additionally, the prefetch capacity limit is now dynamically calculated to adapt to different host and device memory configurations, ensuring efficient memory utilization across varied hardware setups.

Highlights

Relaxed Host Memory Size Constraint: The previous assertion requiring host memory to be strictly larger than device memory has been commented out, enabling configurations where host memory can be smaller than device memory.
Improved Eviction Resilience: The evict method now gracefully handles failures during write_backup operations (e.g., due to host OOM) by falling back to a regular eviction strategy, preventing data loss or system instability.
Enhanced Host Memory Eviction Logic: The evict_host method has been updated to allow eviction of host memory for nodes that are still resident on the device, which is crucial for reclaiming host memory when the host cache is smaller than the device cache.
Dynamic Prefetch Capacity Limit: The calculation for prefetch_capacity_limit now dynamically adjusts based on whether host memory is larger or smaller than device memory, optimizing prefetching behavior for various memory configurations.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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. ↩

stmatengss · 2026-01-11T16:54:37Z

/tag-and-rerun-ci

gemini-code-assist

Code Review

This pull request introduces support for host memory being smaller than device memory in HiCache. The changes are logical and well-implemented. In cache_controller.py, the logic for calculating prefetch_capacity_limit is correctly adapted for this new scenario. In hiradix_cache.py, the evict method is improved for robustness, and the evict_host method is updated to allow evicting host memory for nodes still on the device, which is a key part of this feature. Finally, a now-invalid assertion in memory_pool_host.py is addressed. I have a couple of minor suggestions to improve code quality.

gemini-code-assist · 2026-01-11T16:55:59Z

+            if self.mem_pool_host.size > self.mem_pool_device.size:
+                self.prefetch_capacity_limit = int(
+                    0.8 * (self.mem_pool_host.size - self.mem_pool_device.size)
+                )
+            else:
+                self.prefetch_capacity_limit = int(0.5 * self.mem_pool_host.size)


The magic numbers 0.8 and 0.5 for calculating prefetch_capacity_limit should be defined as named constants at the module level to improve readability and maintainability. For example, PREFETCH_CAPACITY_FACTOR_LARGE_HOST = 0.8 and PREFETCH_CAPACITY_FACTOR_SMALL_HOST = 0.5.

gemini-code-assist · 2026-01-11T16:55:59Z

+        # assert (
+        #     self.size > device_pool.size
+        # ), "The host memory should be larger than the device memory with the current protocol"


Instead of commenting out the assertion, it's better to remove it completely to keep the code clean, since this assertion is no longer valid with the new changes.

let's only skip this check when write back is selected, or when user specified smaller CPU memory size, fall back to write back policy

Sure, we need a check here

stmatengss · 2026-01-18T05:18:43Z

/rerun-failed-ci

xiezhq-hermann · 2026-01-20T03:03:08Z

please take a look at this PR as well:
#14718

Resolved conflicts in: - python/sglang/srt/managers/cache_controller.py: Applied L2 < L1 fix to the new attach_storage_backend method - python/sglang/srt/mem_cache/hiradix_cache.py: Merged evictable_host_leaves tracking with conditional eviction logic Changes preserved from this branch: 1. Support for host memory (L2) smaller than device memory (L1) 2. Conditional prefetch_capacity_limit calculation based on memory sizes 3. Fallback to regular eviction when write_backup fails (returns 0) 4. Conditional tree cleanup only when node is evicted from device 5. Commented out assertion requiring host_size > device_size Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

stmatengss

Review: [HiCache] Support host memory size smaller than device memory size

Key Issue: `evict_host` new logic for non-evicted nodes is unreachable

evict_host builds its heap from self.evictable_host_leaves (line 844). But _update_host_leaf_status (line 759-772) only adds nodes where node.evicted is True — non-evicted nodes are excluded by the early return at line 760:

def _update_host_leaf_status(self, node: TreeNode):
    if not node.evicted or node.lock_ref > 0:
        if node in self.evictable_host_leaves:
            self.evictable_host_leaves.remove(node)
        return

This means backed-up nodes still on device will never enter the eviction heap. The new branches in evict_host that handle x.evicted is False (keeping the node in the tree, only clearing host_value) are dead code.

For the PR's goal to work — freeing redundant host copies of on-device nodes when host < device — _update_host_leaf_status needs to be updated to also consider non-evicted nodes that have host_value is not None.

Minor items

Commented-out assertion in memory_pool_host.py: If the constraint is intentionally lifted, delete the assertion entirely rather than commenting it out.
Magic number 0.5 in prefetch_capacity_limit fallback (cache_controller.py:456): Worth a brief comment explaining the rationale for reserving half of host memory for prefetch when host <= device.
Tree invariant is subtle but correct: After evict_host, a node reachable in the tree with evicted=True will always have host_value != None, because evict_host removes evicted nodes from the tree when clearing their host_value. This invariant protects match_prefix_return_kv (line 1156-1157) from hitting len(None). A comment documenting this invariant would help future readers.

stmatengss · 2026-02-25T17:16:25Z

/rerun-failed-ci

stmatengss · 2026-02-26T02:29:33Z

/rerun-failed-ci

Signed-off-by: Xuchun Shang <xuchun.shang@linux.alibaba.com>

XucSh · 2026-03-04T06:56:40Z

/rerun-failed-ci

XucSh · 2026-03-04T09:51:27Z

/rerun-failed-ci

XucSh · 2026-03-05T01:45:55Z

/rerun-failed-ci

xiezhq-hermann · 2026-03-10T01:22:42Z

            num_evicted += self.cache_controller.evict_host(x.host_value)
+            x.host_value = None
+
+            if x.evicted:


does this mean we can't evict the host if the data is present on GPU?

Data in L2 should be consistent with the GPU.

stmatengss · 2026-03-13T17:43:02Z

#20535 should be implemented more clearly. This PR might not cover complex cases. @xiezhq-hermann

stmatengss · 2026-03-13T17:52:16Z

After discussing with @vladnosiv, I totally agree with a new --hicache-host-memory-mode buffer_only parameter. When L2 is less than L3, it could improve performance by reducing copying.

[HiCache] Support host memory size smaller than device memory size

3d9f693

stmatengss requested review from Ying1123, hanming-lu, hnyls2002, merrymercy, xiezhq-hermann and yizhang2077 as code owners January 11, 2026 16:54

github-actions Bot added the run-ci label Jan 11, 2026

gemini-code-assist Bot reviewed Jan 11, 2026

View reviewed changes

stmatengss commented Jan 11, 2026

View reviewed changes

Comment thread python/sglang/srt/mem_cache/hiradix_cache.py Outdated

Update python/sglang/srt/mem_cache/hiradix_cache.py

cbfb249

stmatengss assigned xiezhq-hermann Jan 12, 2026

xiezhq-hermann added the high priority label Jan 15, 2026

xiezhq-hermann assigned stmatengss Jan 27, 2026

stmatengss and others added 2 commits February 10, 2026 01:28

Merge branch 'main' into mateng/make_l2_can_less_than_l1

2f80b3a

stmatengss commented Feb 25, 2026

View reviewed changes

update

35fd328

github-actions Bot added the hicache Hierarchical Caching for SGLang label Feb 25, 2026

XucSh added 2 commits March 3, 2026 11:40

Merge branch 'main' into mateng/make_l2_can_less_than_l1

cf0af9a

fix lint

f6005cf

Signed-off-by: Xuchun Shang <xuchun.shang@linux.alibaba.com>

xiezhq-hermann reviewed Mar 10, 2026

View reviewed changes

xiezhq-hermann mentioned this pull request Jun 18, 2026

[HiCache] remove large host mem constraint #28614

Merged

5 tasks

xiezhq-hermann closed this Jun 18, 2026

xiezhq-hermann mentioned this pull request Aug 14, 2026

[HiCache] Buffer-only mode for HiCache host memory layer #34798

Merged

5 tasks

Conversation

stmatengss commented Jan 11, 2026

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Review Process

Uh oh!

gemini-code-assist Bot commented Jan 11, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

stmatengss commented Jan 11, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jan 11, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jan 11, 2026

Choose a reason for hiding this comment

Uh oh!

xiezhq-hermann Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

stmatengss Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

stmatengss commented Jan 18, 2026

Uh oh!

xiezhq-hermann commented Jan 20, 2026

Uh oh!

stmatengss left a comment

Choose a reason for hiding this comment

Review: [HiCache] Support host memory size smaller than device memory size

Key Issue: evict_host new logic for non-evicted nodes is unreachable

Minor items

Uh oh!

stmatengss commented Feb 25, 2026

Uh oh!

stmatengss commented Feb 26, 2026

Uh oh!

XucSh commented Mar 4, 2026

Uh oh!

XucSh commented Mar 4, 2026

Uh oh!

XucSh commented Mar 5, 2026

Uh oh!

xiezhq-hermann Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

stmatengss Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

stmatengss commented Mar 13, 2026

Uh oh!

stmatengss commented Mar 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Key Issue: `evict_host` new logic for non-evicted nodes is unreachable