Moe kernels refactor

[Datta0]

Tried to preserve as much authorship as possible with the help from codex.

Ref: https://github.com/unslothai/unsloth/pull/4145/changes

[gemini-code-assist]

Summary of Changes

Hello, 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 significantly enhances the performance of Mixture-of-Experts (MoE) MLP blocks by introducing a new, highly optimized grouped GEMM implementation. By fusing critical operations directly into Triton kernels, the changes aim to minimize computational overheads and memory transfers. The integration includes an intelligent autotuning and caching mechanism to dynamically adapt and store optimal kernel configurations, ensuring peak efficiency. This refactor provides a more performant and maintainable foundation for MoE architectures.

Highlights

• Optimized MoE Kernel Implementation: Introduced a new, highly optimized grouped GEMM implementation for Mixture-of-Experts (MoE) MLP blocks, leveraging Triton kernels for enhanced performance.
• Fused Operations: Implemented fusions for key MoE operations, including permutation of hidden states (token to expert order and vice-versa) and multiplication by expert weights, directly within the GEMM kernels to reduce overhead.
• Autotuning and Caching System: Added a robust system for autotuning Triton kernel configurations and caching them to disk, ensuring optimal performance across different hardware setups without repetitive tuning.
• Comprehensive Testing and Benchmarking: Included extensive unit and end-to-end tests, along with benchmarking scripts, to validate the correctness and measure the performance improvements of the new MoE grouped GEMM kernels for Qwen3 and Llama4 models.
• Licensing Update: All newly added kernel files are licensed under the GNU Affero General Public License v3.0, ensuring open and collaborative use.
• Temporary Patch Updates: Updated existing temporary patches to correctly reference the new MoE kernel module structure and ensure compatibility with LoRA extractor functions.

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

Changelog

• unsloth_zoo/kernels/init.py
  • Added package namespace.
• unsloth_zoo/kernels/moe/LICENSE
  • Added GNU Affero General Public License v3.0.
• unsloth_zoo/kernels/moe/README.md
  • Added detailed documentation for MoE Grouped GEMM, including background, optimizations, structure, tests, benchmarks, and notes.
• unsloth_zoo/kernels/moe/init.py
  • Added package namespace with AGPLv3 license header.
• unsloth_zoo/kernels/moe/autotune_cache.py
  • Added a system for caching and retrieving autotuned MoE kernel configurations.
• unsloth_zoo/kernels/moe/benchmark/benchmark_fused_moe.py
  • Added a script for benchmarking fused MoE implementations against Hugging Face models.
• unsloth_zoo/kernels/moe/benchmark/utils.py
  • Added utility functions for creating, processing, and saving benchmark results.
• unsloth_zoo/kernels/moe/grouped_gemm/LICENSE
  • Added GNU Affero General Public License v3.0.
• unsloth_zoo/kernels/moe/grouped_gemm/init.py
  • Added package namespace with AGPLv3 license header.
• unsloth_zoo/kernels/moe/grouped_gemm/interface.py
  • Added Python interface for grouped GEMM operations, including forward and backward passes, and TMA support checks.
• unsloth_zoo/kernels/moe/grouped_gemm/kernels/init.py
  • Added package namespace with AGPLv3 license header.
• unsloth_zoo/kernels/moe/grouped_gemm/kernels/autotuning.py
  • Added utilities for autotuning Triton kernel configurations, including default block sizes and pruning criteria.
• unsloth_zoo/kernels/moe/grouped_gemm/kernels/backward.py
  • Added Triton kernels for grouped GEMM backward passes (dX and dW), including autotuning.
• unsloth_zoo/kernels/moe/grouped_gemm/kernels/forward.py
  • Added Triton kernel for grouped GEMM forward pass, including autotuning and fusions.
• unsloth_zoo/kernels/moe/grouped_gemm/kernels/tuning.py
  • Added data classes for kernel configurations and results, along with pruning functions for manual tuning.
• unsloth_zoo/kernels/moe/grouped_gemm/reference/init.py
  • Added package namespace with AGPLv3 license header.
• unsloth_zoo/kernels/moe/grouped_gemm/reference/layers/llama4_moe.py
  • Added reference implementations for Llama4 MoE block using both torch-native and Triton grouped GEMM.
• unsloth_zoo/kernels/moe/grouped_gemm/reference/layers/qwen3_moe.py
  • Added reference implementations for Qwen3 MoE block using both torch-native and Triton grouped GEMM.
• unsloth_zoo/kernels/moe/grouped_gemm/reference/moe_block.py
  • Added a reference implementation of Qwen3 MoE block using Triton grouped GEMM.
• unsloth_zoo/kernels/moe/grouped_gemm/reference/moe_ops.py
  • Added utility functions for MoE operations like permuting, unpermuting, calculating top-k, and torch-native grouped GEMM.
• unsloth_zoo/kernels/moe/requirements.txt
  • Added new Python dependencies for the MoE kernels.
• unsloth_zoo/kernels/moe/tests/init.py
  • Added package namespace with AGPLv3 license header.
• unsloth_zoo/kernels/moe/tests/common.py
  • Added common utilities and data structures for MoE kernel testing.
• unsloth_zoo/kernels/moe/tests/moe_utils.py
  • Added utility functions for testing MoE blocks, including gradient checking and result comparison.
• unsloth_zoo/kernels/moe/tests/run_qwen3_moe_tests.sh
  • Added a shell script to run Qwen3 MoE tests with various configurations.
• unsloth_zoo/kernels/moe/tests/test_grouped_gemm.py
  • Added comprehensive unit tests for grouped GEMM forward and backward kernels.
• unsloth_zoo/kernels/moe/tests/test_llama4_moe.py
  • Added tests for Llama4 MoE integration with the new grouped GEMM kernels.
• unsloth_zoo/kernels/moe/tests/test_qwen3_moe.py
  • Added tests for Qwen3 MoE integration with the new grouped GEMM kernels.
• unsloth_zoo/temporary_patches/moe_utils.py
  • Updated import paths for grouped_gemm and autotune_cache.
  • Added staticmethod decorator to _unsloth_lora_extractor_fn.
• unsloth_zoo/temporary_patches/qwen3_5_moe.py
  • Applied staticmethod decorator to _unsloth_lora_extractor_fn.
• unsloth_zoo/temporary_patches/qwen3_moe.py
  • Applied staticmethod decorator to _unsloth_lora_extractor_fn.
• unsloth_zoo/temporary_patches/qwen3_next_moe.py
  • Applied staticmethod decorator to _unsloth_lora_extractor_fn.
• unsloth_zoo/temporary_patches/qwen3_vl_moe.py
  • Applied staticmethod decorator to _qwen3_vl_lora_extractor.

Activity

• The pull request introduces a new moe kernel package within unsloth_zoo/kernels, containing a complete implementation of optimized grouped GEMM operations for Mixture-of-Experts (MoE) models.
• This includes Triton-based kernels for forward and backward passes, an autotuning and caching system for kernel configurations, and extensive testing and benchmarking infrastructure.
• The author, Datta0, noted efforts to preserve authorship, suggesting the use of tools like Codex for code generation while maintaining proper attribution.
• Several existing temporary patches related to Qwen3 and other MoE models were updated to correctly integrate with the new kernel structure and ensure proper function calls.

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

This pull request introduces a significant refactoring by moving Mixture-of-Experts (MoE) kernels into a new unsloth_zoo package. This includes new, optimized Triton kernels for grouped GEMM operations, along with infrastructure for autotuning and caching. The changes are extensive and well-structured. My review focuses on the correctness, performance, and maintainability of this new kernel infrastructure. I've identified a few issues, including a bug that could lead to a runtime error, some performance inefficiencies in memory allocation for Triton kernels, and some dead code. Addressing these points will help improve the robustness and performance of the new MoE kernels.