add qwen3 attn

[poryfly]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @poryfly, 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 integrates dedicated attention mechanisms and a comprehensive set of optimization rules for Qwen3Moe models within the ktransformers library. The primary goal is to enhance the performance and efficiency of Qwen3Moe models by replacing standard PyTorch modules with specialized ktransformers operators, with a particular emphasis on leveraging hardware acceleration like Intel AMX for key components such as attention, linear layers, and Mixture-of-Experts blocks.

Highlights

• Qwen3Moe Attention Implementation: A new KQwen3MoeAttention class has been introduced, providing an optimized attention mechanism specifically for Qwen3Moe models. This implementation includes the application of Rotary Position Embeddings (RoPE) using Qwen3MoeRotaryEmbedding and dynamically selects between eager, SDPA, or FlashAttention-2 implementations based on configuration.
• Qwen3Moe Optimization Rules: A new YAML configuration file, Qwen3Moe-sft-amx.yaml, has been added. This file defines a comprehensive set of optimization rules for Qwen3Moe models, replacing standard PyTorch modules with ktransformers optimized operators for various components like linear layers, rotary embeddings, Mixture-of-Experts (MoE) blocks, and RMS normalization.
• AMXInt8 Backend Support for MoE: The newly added optimization rules explicitly configure the Mixture-of-Experts (MoE) experts to utilize an AMXInt8 backend. This indicates a strategic focus on leveraging Intel Advanced Matrix Extensions for performance acceleration through quantized integer operations, particularly for the computationally intensive MoE layers.

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 adds support for Qwen3 attention by introducing a new KQwen3MoeAttention operator and an optimization rule file Qwen3Moe-sft-amx.yaml. The changes are generally in the right direction, but I've identified a few critical issues that need to be addressed.

In attention.py, the attention implementation is hardcoded, which limits flexibility.
In Qwen3Moe-sft-amx.yaml, there are two significant problems:

1. An incorrect class is targeted for RoPE optimization, which will cause it to fail.
2. A model replacement rule is duplicated and points to an incorrect model class for Qwen2 instead of Qwen3.

These issues could lead to incorrect model behavior or failed optimizations. Please see the detailed comments for suggestions on how to fix them.

[gemini-code-assist]

The match rule for RotaryEmbedding replacement is targeting ktransformers.models.modeling_qwen2_moe.Qwen2MoeRotaryEmbedding. However, the Qwen3 model uses ktransformers.models.modeling_qwen3_moe.Qwen3MoeRotaryEmbedding. This will cause the RoPE optimization to be skipped for Qwen3 models. Please correct the class path.

    class: ktransformers.models.modeling_qwen3_moe.Qwen3MoeRotaryEmbedding

[gemini-code-assist]

The attention implementation is hardcoded to "flash_attention_2". This overrides any configuration and prevents the use of other attention mechanisms like "eager" or "sdpa". This line should be removed to allow the attention implementation to be controlled by the model's configuration.

[gemini-code-assist]

There are two match rules for name: "^model$". The first one incorrectly specifies class: "ktransformers.operators.models.KQwen2MoeModel", which is for Qwen2 models, not Qwen3. This is likely a copy-paste error and is redundant with the correct rule at the end of the file. Please remove this incorrect and duplicated rule block.

[KMSorSMS]

Please add some description to this PR.

[JimmyPeilinLi]

Please add some description to this PR.

Wait for a moment. We are working together to improve it.