add Step-3.5-Flash model support

[yhyang201]

Motivation

add Step-3.5-Flash model support

Important notes:
The current implementation differs slightly from the standard MTP implementation in Step3.5 Flash. In the standard multi-layer MTP design of Step3.5 Flash, the hidden states of each MTP layer are passed from the preceding MTP layer (the hidden states of the initial (layer-0) MTP still being provided by the target model). In contrast, the current SGL implementation obtains hidden states directly from the target model for all MTP layers. Empirical evaluations indicate that the overall performance remains strong; however, this design choice may lead to a slight reduction in acceptance rate in certain scenarios. This behavior will be corrected shortly, and we expect to implement the standard multi-layer MTP design of Step3.5 Flash in the near future.

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

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 @yhyang201, 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 full support for the Step3.5 Flash model, encompassing its base architecture and a specialized Multi-Token Predictor (MTP) variant. It includes new configuration definitions, detailed implementations of attention and Mixture-of-Experts (MoE) layers, and integrates the model with speculative decoding capabilities. The changes ensure the new model can be properly configured, loaded, and utilized within the system, with a noted temporary design choice in MTP hidden state handling that will be addressed in future updates.

Highlights

• New Model Support: Introduced comprehensive support for the Step3.5 Flash model architecture, including its base and Multi-Token Predictor (MTP) variants.
• Configuration and Model Definitions: Added a dedicated Step3p5Config for model parameters and implemented core components like Step3p5MLP, Step3p5MoEMLP, Step3p5Attention, and Step3p5DecoderLayer.
• Mixture-of-Experts (MoE) and Attention Enhancements: Integrated specific logic for handling MoE layers, including routing and expert execution, and refined attention mechanisms with rotary embeddings and RadixAttention, supporting sliding and full attention types.
• Multi-Token Predictor (MTP) Implementation: Developed Step3p5MTP for efficient multi-token prediction, noting a current deviation from standard MTP where hidden states are directly obtained from the target model for all MTP layers, with plans for future alignment.
• Speculative Decoding Integration: Enabled multi-layer EAGLE speculative decoding for Step3p5NextForCausalLM and adjusted hierarchical cache settings for improved performance.
• System Integration: Updated model configuration, function call parsing, and reasoning parsing to recognize and properly handle the new Step3.5 models.

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

[yhyang201]

/tag-and-rerun-ci

[gemini-code-assist]

Code Review

This pull request introduces support for the Step3.5 flash model and its MTP variant, which is a significant feature addition. The changes are comprehensive, touching upon model configurations, implementations, and parsers. While the overall implementation appears solid, I've identified a few areas for improvement, including a potential runtime error due to a type mismatch in the model configuration, a possible TypeError from an unsafe dictionary access, and some opportunities to improve code clarity by removing redundancy and hardcoded values.

[yhyang201]

/rerun-failed-ci