[RadixTree][8/N Refactor]: unify lock interface by pansicheng · Pull Request #20330 · sgl-project/sglang

pansicheng · 2026-03-11T07:45:36Z

Motivation

The current lock reference interface is inconsistent across cache implementations:

inc_lock_ref returns different types depending on the implementation.
dec_lock_ref uses inconsistent parameter types and calling conventions.
dec_lock_ref also returns different types across implementations.
These inconsistencies make the interface harder to understand, use, and maintain.

This PR attempts to unify the lock interface across all cache implementations by making the behavior of inc_lock_ref and dec_lock_ref consistent.

Modifications

Accuracy Tests

python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319

# full: qwen3-8b
Accuracy: 0.905
Invalid: 0.000

# mamba: qwen3.5-9b
Accuracy: 0.898
Invalid: 0.000

# swa: gpt-oss-20b (--enable-thinking)
Accuracy: 0.862
Invalid: 0.001

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-03-11T07:46:34Z

Summary of Changes

Hello, 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 refactors the lock reference interface within the cache implementations to address existing inconsistencies. By introducing dedicated dataclasses for the input parameters and return values of inc_lock_ref and dec_lock_ref, the change aims to create a more uniform, predictable, and maintainable API for managing cache locks across different cache types. This standardization simplifies future development and reduces potential errors arising from varying interface behaviors.

Highlights

Unified Lock Interface: Standardized the inc_lock_ref and dec_lock_ref methods across all cache implementations to ensure consistent behavior and return types.
New Dataclasses for Lock Operations: Introduced IncLockRefResult, DecLockRefParams, and DecLockRefResult dataclasses to define explicit return types and parameters for lock reference operations, improving type safety and clarity.
Interface Adaptation: Updated all existing calls and implementations of inc_lock_ref and dec_lock_ref in various cache modules to conform to the new standardized interface.

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

Changelog

python/sglang/srt/managers/schedule_policy.py
- Imported DecLockRefParams for standardized lock reference decrement parameters.
- Updated _req_inc_lock_ref to use the IncLockRefResult from inc_lock_ref.
- Modified _lock_node to use IncLockRefResult and DecLockRefParams for lock reference operations.
python/sglang/srt/mem_cache/base_prefix_cache.py
- Added IncLockRefResult dataclass to standardize the return type of inc_lock_ref.
- Added DecLockRefParams dataclass to standardize parameters for dec_lock_ref.
- Added DecLockRefResult dataclass to standardize the return type of dec_lock_ref.
- Updated abstract method signatures for inc_lock_ref and dec_lock_ref to use the new dataclasses.
python/sglang/srt/mem_cache/chunk_cache.py
- Imported DecLockRefParams, DecLockRefResult, and IncLockRefResult.
- Updated inc_lock_ref to return IncLockRefResult.
- Updated dec_lock_ref to accept DecLockRefParams and return DecLockRefResult.
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
- Imported DecLockRefParams, DecLockRefResult, and IncLockRefResult.
- Updated load_back to extract delta from the IncLockRefResult.
- Updated inc_lock_ref to return IncLockRefResult.
- Updated dec_lock_ref to accept DecLockRefParams and return DecLockRefResult.
python/sglang/srt/mem_cache/hiradix_cache.py
- Imported DecLockRefParams, DecLockRefResult, and IncLockRefResult.
- Updated inc_lock_ref to return IncLockRefResult.
- Updated dec_lock_ref to accept DecLockRefParams and return DecLockRefResult.
- Modified load_back to extract delta from the IncLockRefResult.
python/sglang/srt/mem_cache/mamba_radix_cache.py
- Imported DecLockRefParams, DecLockRefResult, and IncLockRefResult.
- Updated inc_lock_ref to return IncLockRefResult.
- Updated dec_lock_ref to accept DecLockRefParams and return DecLockRefResult.
python/sglang/srt/mem_cache/radix_cache.py
- Imported DecLockRefParams, DecLockRefResult, and IncLockRefResult.
- Updated inc_lock_ref to return IncLockRefResult.
- Updated dec_lock_ref to accept DecLockRefParams and return DecLockRefResult.
python/sglang/srt/mem_cache/radix_cache_cpp.py
- Imported Optional, DecLockRefParams, DecLockRefResult, and IncLockRefResult.
- Updated dec_lock_ref to accept DecLockRefParams and return DecLockRefResult.
- Updated inc_lock_ref to return IncLockRefResult.
python/sglang/srt/mem_cache/session_aware_cache.py
- Imported DecLockRefParams, DecLockRefResult, and IncLockRefResult.
- Updated inc_lock_ref to return IncLockRefResult.
- Updated dec_lock_ref to accept DecLockRefParams and return DecLockRefResult.
- Modified release_session to use DecLockRefParams when calling dec_lock_ref.
python/sglang/srt/mem_cache/swa_radix_cache.py
- Imported DecLockRefParams, DecLockRefResult, and IncLockRefResult.
- Modified cache_finished_req to use DecLockRefParams when calling dec_lock_ref.
- Updated cache_unfinished_req to use DecLockRefParams and extract swa_uuid_for_lock from IncLockRefResult.
- Updated inc_lock_ref to return IncLockRefResult.
- Updated dec_lock_ref to accept DecLockRefParams and return DecLockRefResult.
test/registered/scheduler/test_prefill_adder.py
- Imported DecLockRefResult and IncLockRefResult for testing purposes.
- Updated mock return values for inc_lock_ref and dec_lock_ref to reflect the new dataclass types.

Activity

The pull request includes accuracy test results for qwen3-8b, qwen3.5-9b, and gpt-oss-20b models, showing consistent accuracy after the changes.
The author provided a detailed motivation for the refactoring, highlighting inconsistencies in the lock reference interface.

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

hzh0425 · 2026-03-11T07:55:57Z

/tag-and-rerun-ci

gemini-code-assist

Code Review

This pull request successfully unifies the lock interface (inc_lock_ref and dec_lock_ref) across various cache implementations by introducing IncLockRefResult, DecLockRefParams, and DecLockRefResult data classes. The changes are well-structured and improve the consistency and maintainability of the codebase.

I've identified a critical bug in swa_radix_cache.py that could lead to a crash due to a potential AttributeError when handling None parameters. I've also found a minor issue in the test setup in test_prefill_adder.py where a mock is returning a class instead of an instance.

Please address these points to ensure the stability and correctness of the changes.

_{Note: Security Review did not run due to the size of the PR.}

pansicheng · 2026-03-13T02:16:54Z

/rerun-failed-ci

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

pansicheng · 2026-03-13T05:46:52Z

/rerun-failed-ci

pansicheng requested review from Ying1123, hanming-lu, hnyls2002, hzh0425, ispobock, merrymercy, xiezhq-hermann and yizhang2077 as code owners March 11, 2026 07:45

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

gemini-code-assist Bot reviewed Mar 11, 2026

View reviewed changes

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

Comment thread test/registered/scheduler/test_prefill_adder.py Outdated

hzh0425 self-assigned this Mar 11, 2026

ispobock mentioned this pull request Mar 12, 2026

[Roadmap] Unified Hybrid Radix Cache Refactor #20415

Open

23 tasks

ispobock approved these changes Mar 12, 2026

View reviewed changes

hzh0425 approved these changes Mar 12, 2026

View reviewed changes

xiezhq-hermann approved these changes Mar 12, 2026

View reviewed changes

pansicheng and others added 2 commits March 13, 2026 02:25

unified lock interface

b929bcd

Apply suggestions from code review

fc79c06

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

pansicheng force-pushed the unify_incdec_lockref branch from 878b2db to fc79c06 Compare March 13, 2026 02:26

ispobock added 2 commits March 14, 2026 18:41

Merge branch 'main' into unify_incdec_lockref

d6c9942

Merge branch 'main' into unify_incdec_lockref

c016a4f

ispobock merged commit 97b2a89 into sgl-project:main Mar 16, 2026
134 of 146 checks passed

Wangzheee pushed a commit to Wangzheee/sglang that referenced this pull request Mar 21, 2026

[RadixTree][8/N Refactor]: unify lock interface (sgl-project#20330)

1cd73ca

0-693 pushed a commit to 0-693/sglang that referenced this pull request Mar 25, 2026

[RadixTree][8/N Refactor]: unify lock interface (sgl-project#20330)

d49dc40

JustinTong0323 pushed a commit to JustinTong0323/sglang that referenced this pull request Apr 7, 2026

[RadixTree][8/N Refactor]: unify lock interface (sgl-project#20330)

f98a216

yhyang201 pushed a commit to yhyang201/sglang that referenced this pull request Apr 22, 2026

[RadixTree][8/N Refactor]: unify lock interface (sgl-project#20330)

67b297d

hnyls2002 mentioned this pull request Apr 28, 2026

Deepseek V4 #23882

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[RadixTree][8/N Refactor]: unify lock interface#20330

[RadixTree][8/N Refactor]: unify lock interface#20330
ispobock merged 4 commits intosgl-project:mainfrom
pansicheng:unify_incdec_lockref

pansicheng commented Mar 11, 2026

Uh oh!

gemini-code-assist Bot commented Mar 11, 2026

Uh oh!

hzh0425 commented Mar 11, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

pansicheng commented Mar 13, 2026

Uh oh!

pansicheng commented Mar 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

pansicheng commented Mar 11, 2026

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Review Process

Uh oh!

gemini-code-assist Bot commented Mar 11, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

hzh0425 commented Mar 11, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

pansicheng commented Mar 13, 2026

Uh oh!

pansicheng commented Mar 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants