[CPU]Improve cpu fused moe perf #27244
Conversation
- Add oneDNN/ACL matmul path for AArch64 - Use silu_and_mul Op Signed-off-by: Zhang Xiangze <[email protected]>
![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 improves CPU performance for fused MoE layers by introducing a oneDNN/ACL matmul path for AArch64 architectures and leveraging a custom silu_and_mul operator. The changes are well-implemented, preparing oneDNN matmul handlers during initialization and using them in the forward pass. While this is a solid performance enhancement, I've identified a high-severity issue concerning model serialization. The created oneDNN handlers are not serializable, and attaching them to the model can lead to crashes upon saving and reloading the model.
| gate_up_handle = ops.create_onednn_mm(layer_w13_weight.t(), 32) | ||
| layer.gate_up_linear.append( | ||
| lambda x, handle=gate_up_handle, bias=layer_w13_bias: ops.onednn_mm( | ||
| handle, x, bias | ||
| ) | ||
| ) | ||
| down_handle = ops.create_onednn_mm(layer_w2_weight.t(), 32) | ||
| layer.down_linear.append( | ||
| lambda x, handle=down_handle, bias=layer_w2_bias: ops.onednn_mm( | ||
| handle, x, bias | ||
| ) | ||
| ) |
There was a problem hiding this comment.
The CPUDNNLGEMMHandler objects created by ops.create_onednn_mm contain pointers to C++ state and are not serializable. Storing lambdas that capture these handlers in layer.gate_up_linear and layer.down_linear will cause issues if the model is serialized (e.g., with pickle or torch.save). Upon deserialization, the handler pointers will be invalid, which can lead to segmentation faults when the model is used or garbage collected.
To prevent this, the CPUDNNLGEMMHandler class should be made non-picklable by implementing __getstate__ to raise an exception. Since that class is not in this file, an alternative is to avoid storing these handlers on the torch.nn.Module instance if model serialization is a possibility. If serialization is not a supported use case for CPU-based models, this might be acceptable, but it's a significant risk.
![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.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Signed-off-by: Zhang Xiangze <[email protected]>
|
I saw vllm crash when running unit test locally with It seems the issue is not directly related with this PR. But the moe test exposed an issue in onednn MatMul cache. The issue: The tests passed when I tested bf16 and fp32 separately. |
|
@mgoin @bigPYJ1151 Can you help to review this PR? |
|
Hi @xiangze-arm - can you test again with #27472? |
3 participants
Tested locally with Qwen3-30B-A3B model, with this PR, MoE can go into onednn_mm path on AArch64 CPU.