[RadixTree][1/N Refactor]: Support unified match_prefix params

[hzh0425]

Motivation

Modifications

Accuracy Tests

# mamba
Accuracy: 0.950
Invalid: 0.000
Latency: 14.648 s
Output throughput: 5448.861 token/s


# llama4
Accuracy: 0.936
Invalid: 0.000
Latency: 35.233 s
Output throughput: 1452.040 token/s

# qwen32b
Accuracy: 0.940
Invalid: 0.000
Latency: 25.780 s
Output throughput: 2833.223 token/s

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

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

[gemini-code-assist]

Summary of Changes

Hello @hzh0425, 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 match_prefix function within the RadixTree caching system by introducing a unified parameter object, MatchPrefixParams. This change streamlines the API for prefix matching operations, making it more consistent and easier to extend, particularly for specialized cache types like Mamba. The update involves modifying the method signatures and all call sites to use this new dataclass, enhancing code readability and maintainability across the caching infrastructure.

Highlights

• Unified match_prefix Parameters: Introduced a new MatchPrefixParams dataclass to standardize the arguments passed to the match_prefix method across various cache implementations, replacing the previous key and **kwargs pattern.
• API Modernization: The match_prefix method signature in BasePrefixCache and all its derived classes (e.g., RadixCache, MambaRadixCache, LmcRadixCache) has been updated to accept a single MatchPrefixParams object, improving clarity and maintainability.
• Mamba-Specific Parameter Encapsulation: Mamba-specific parameters like req and cow_mamba are now explicitly part of the MatchPrefixParams dataclass, centralizing their handling and making the API more robust.
• Extensive Call Site Updates: All existing calls to match_prefix throughout the codebase, including in scheduling managers and various cache implementations, have been updated to conform to the new MatchPrefixParams API.
• Test Suite Alignment: Unit tests for MambaRadixCache, RadixCache, and SwaRadixCache have been modified to use the new MatchPrefixParams object when calling match_prefix, ensuring continued correctness.

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

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

[hzh0425]

/tag-and-rerun-ci

[gemini-code-assist]

Code Review

This pull request refactors the match_prefix method across various cache implementations to use a unified MatchPrefixParams dataclass instead of a flexible **kwargs signature. This is a great improvement for code clarity and maintainability. The changes have been applied consistently across all relevant files, including base classes, implementations, and tests. I've found one minor area for improvement to avoid a redundant function call, but overall the refactoring is well-executed.