Fix the performance degradation issue of the single-wheel operation in Ant Moving.

[oagniqgnat]

Fix the performance degradation issue of the single-wheel operation in Ant Moving.

[gemini-code-assist]

Summary of Changes

Hello @oagniqgnat, 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 addresses a performance degradation observed during single-wheel operations within the Ant Moving system, specifically impacting the Mixture of Experts (MoE) combine normal operator. The solution involves a significant architectural refactoring to differentiate between single-round and multi-round processing paths. By isolating the more complex multi-round logic into a dedicated implementation and optimizing the original operator for single-round execution, the system avoids unnecessary overhead. Complementary memory optimizations and conditional dispatching further contribute to the overall performance enhancement for the targeted single-wheel scenarios.

Highlights

• MoE Combine Operation Refactoring: The CamMoeCombineNormal operator has been refactored to separate single-round and multi-round processing logic into distinct implementations, improving clarity and maintainability.
• Single-Round Performance Optimization: The original CamMoeCombineNormal implementation is now streamlined for single-round operations by removing multi-round specific overhead, leading to performance improvements.
• Dedicated Multi-Round Implementation: A new CamMoeCombineNormalMultiRound class and header file (cam_moe_combine_normal_multi_round.h) have been introduced to encapsulate the complex multi-round processing logic, which is dynamically selected at runtime.
• Memory Footprint Reduction: The COMBINE_STATE_WIN_OFFSET constant, a key memory buffer offset, has been reduced from 8MB to 4MB across several files, contributing to memory optimization.
• Conditional Dispatch Logic: Updates to notify_dispatch.h and cam_moe_dispatch_normal.h ensure that multi-round specific synchronization and data handling are only engaged when actually required (i.e., round > 1), further enhancing efficiency for single-round scenarios.

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

[gemini-code-assist]

Code Review

The pull request refactors the CamMoeCombineNormal operator into single-round and multi-round implementations, introducing a tilingKey mechanism to select between them. The original cam_moe_combine_normal.h is streamlined for single-round processing by removing all multi-round specific logic, member variables, and related helper functions, while a new header cam_moe_combine_normal_multi_round.h is added to encapsulate the multi-round functionality. Memory offsets like COMBINE_STATE_WIN_OFFSET are reduced across several files. Review comments highlight issues with tensor management in CopyBufferToShareAndSetStatus, suggesting Get<uint32_t>() instead of AllocTensor<uint32_t>() for TBuf and proper passing of the status tensor to SetStatusBySrcInfo. Additionally, the use of preprocessor macros for TILINGKEY constants is flagged, recommending constexpr variables within a namespace for type safety, and code duplication between the single-round and multi-round classes is noted, suggesting refactoring into a base class. Finally, a redundant using namespace Moe; statement in cam_moe_combine_normal.h is pointed out.