[Feature] Enable return routed experts

[ocss884]

Motivation

As per the request from the RL community, this PR enables sglang to return routed experts (topk) during fwd for later usage in training phase. Thanks the MiMo team for proprosing R3 to help stablizing MoE Reinforcement Learning. This method has been used in MiMo-V2-Flash and DeepSeek-V3.2. We also notice that this methods effectively reduce the logprob diff between training and inference output.

TODO LIST

• enable cuda graph
• setup experts capture buffer
• Test on dpskV3
• Add unit test
• use tobytes to improve performance

This PR does not compatible with:

• deepep usage (will be in next pr)
• mtp (will be in next pr)
• PP (no plan so far)

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.

[gemini-code-assist]

Summary of Changes

Hello @ocss884, 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 implements a new feature that allows the system to return the IDs of the experts chosen by the router in Mixture-of-Experts (MoE) models. This functionality is particularly valuable for the reinforcement learning community, as it provides crucial insights into expert utilization during the training phase. The changes involve creating a dedicated component for capturing these expert IDs, integrating this capture into the model's forward pass, and extending the request and response data structures to propagate this information back to the user.

Highlights

• New Feature: Return Routed Experts: Introduced the capability to return the IDs of routed experts for Mixture-of-Experts (MoE) models, a feature specifically requested by the RL community for training purposes.
• Routed Experts Capturer: Added a new RoutedExpertsCapturer abstraction and its concrete implementation to manage the capture and storage of routed expert IDs during model inference.
• Configuration and API Extension: A new server argument --enable-return-routed-experts and a corresponding enable_return_routed_experts field in ServerArgs have been added to control this feature. Request objects (GenerateReqInput, TokenizedGenerateReqInput) and output structures (BatchTokenIDOutput, BatchStrOutput) are extended to support this new data.
• Integration into Inference Pipeline: The capturing mechanism is integrated into the MoE layer's forward pass, and the captured experts are processed and returned through the scheduler and output processing managers.

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

[ocss884]

@zhaochenyang20 @zhuzilin @yizhang2077 @fzyzcjy

[gemini-code-assist]

Code Review

This pull request introduces a feature to enable returning routed experts for Mixture-of-Experts models, which is useful for training and analysis. The implementation adds a RoutedExpertsCapturer and integrates it throughout the request lifecycle. While the overall approach is sound, there are several critical and high-severity issues that need to be addressed. These include a method signature mismatch that will cause a TypeError, and incorrect logic for clearing the expert data buffer which could lead to data corruption between batches. Additionally, there are opportunities to improve the code quality by removing dead code, a debug print statement, and refactoring the use of global state for better maintainability.

[zhaochenyang20]

🐂🍺

[fzyzcjy]

seems related: #9499

given this todo list, wondering whether reusing EPLB distribution recorder like in 9499 may be good since cudagraph / ep is already supported in it.

(maybe you can collaborate w/ people in 9499)

[fzyzcjy]

offline synced w/ @ocss884, notes

[lizipao]

Hi, I tested your PR and found that there are only output routed_experts, but no input routed_experts. I believe that when doing RL, input routed_experts should be needed. Additionally, if this functionality needs to be implemented, we might need to consider sglang's radix_cache.

[lizipao]

_experts_capturer_host_buffer is insufficient. It is recommended to change it to dynamic expansion or directly modify it to allow custom configuration via parameters?

[yizhang2077]

/rerun-failed-ci try again

[zhaochenyang20]

/rerun-failed-ci

[hnyls2002]

https://github.com/sgl-project/sglang/actions/runs/20393337803?pr=12162

All CUDA CIs passed.

[zhaochenyang20]

GREAT JOB

[merrymercy]

Why do we need to use base64 to encode bytes into string? Can you just use bytes?