Fix MoE routing and TMA checks for PR 3812#3953
Conversation
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
This reverts commit da9e2c4.
This reverts commit 9bb0db2.
* remote unused hybridcache * cleanup
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
Summary of ChangesHello @danielhanchen, I'm Gemini Code Assist1! 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 significantly enhances the performance and robustness of Mixture of Experts (MoE) models, particularly for Qwen3 and GLM4 architectures. It introduces an intelligent caching mechanism for Triton kernel autotuning, refines MoE routing logic, and improves TMA support within grouped GEMM operations. Furthermore, it streamlines the integration of LoRA with MoE models by automating the detection and targeting of expert parameters, and includes a critical fix for vLLM compatibility in GRPO trainers. Highlights
🧠 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 AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces several significant improvements and fixes related to Mixture-of-Experts (MoE) models. Key changes include fixing the MoE routing for Qwen3, making Tensor Memory Access (TMA) checks device-aware, and adding an autotuning cache system for MoE Triton kernels to improve performance by avoiding repeated tuning. The PR also adds automatic LoRA support for MoE expert layers, which simplifies fine-tuning for users. Furthermore, it introduces optimized implementations for GLM-4.7 MoE and a Triton-specific path for Qwen3 MoE. There are also some compatibility fixes for reinforcement learning trainers. Overall, the changes are well-implemented and enhance both performance and user experience. I've found one minor issue with a duplicated line of code.
|
|
||
| return config_fwd, config_bwd_dx, config_bwd_dw | ||
|
|
||
| return config_fwd, config_bwd_dx, config_bwd_dw |
There was a problem hiding this comment.
This return statement is a duplicate of the one on line 371 and is unreachable. It can be safely removed to clean up the code.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b4ee9a4f37
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _SUPPORTS_TMA = {} | ||
|
|
||
| def supports_tma(): | ||
| # Check if triton.set_allocator is available (Triton 3.0+) | ||
| _HAS_SET_ALLOCATOR = hasattr(triton, "set_allocator") |
There was a problem hiding this comment.
_FUSED_MUL_WARN is no longer initialized at module scope, but grouped_gemm_forward still does global _FUSED_MUL_WARN and reads it when fuse_mul_post is True. That path now raises NameError before any kernel launches, so any inference call that enables fused-mul will crash. Reintroduce the module-level flag (e.g., _FUSED_MUL_WARN = False) or guard the access.
Useful? React with 👍 / 👎.
| config_key = (num_experts, hidden_dim, intermediate_dim, top_k, str(dtype)) | ||
|
|
||
| if config_key not in _moe_kernel_configs or force_autotune: |
There was a problem hiding this comment.
get_moe_kernel_configs caches configs by (num_experts, hidden_dim, intermediate_dim, top_k, dtype) only, but get_or_autotune_moe_kernels is device-capability specific. In a mixed-GPU or device-switching process, this cache can reuse an H100 TMA config on an A100 (or vice versa), leading to unsupported TMA paths and runtime failures. Add device index/capability to config_key to keep configs device-safe.
Useful? React with 👍 / 👎.
|
Closing this as it was handled via #3812 |
3 participants
Summary
Testing