Add NPU support for the fused_norm_gate operator #719
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the ✨ Finishing touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello @iiiiLllllzx, 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 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 adds NPU support for the fused_norm_gate operator. While the initiative is great, there are a few critical issues that need to be addressed.
First, the new file fla/modules/fused_norm_gate_npu.py is almost a complete duplicate of fla/modules/fused_norm_gate.py, which poses a significant maintainability problem. This duplication should be avoided by introducing device-specific configurations within the original file.
Second, there's a critical bug in fla/modules/fused_norm_gate_npu.py due to a missing import for autotune_cache_kwargs, which will cause a NameError.
Lastly, the changes in fla/utils.py introduce a regression by incorrectly implementing get_multiprocessor_count, which would cripple performance on non-NPU backends like CUDA. It also re-defines an existing function.
Please see the detailed comments for suggestions on how to fix these issues.
| import triton | ||
| import triton.language as tl | ||
|
|
||
| from fla.utils import get_multiprocessor_count, input_guard |
There was a problem hiding this comment.
The variable autotune_cache_kwargs is used in this file (e.g., on line 28), but it is not imported. This will lead to a NameError at runtime. Please add it to the import statement.
| from fla.utils import get_multiprocessor_count, input_guard | |
| from fla.utils import autotune_cache_kwargs, get_multiprocessor_count, input_guard |
| def get_multiprocessor_count(tensor_idx: int = 0) -> int: | ||
| if triton.runtime.driver.active.get_current_target().backend == 'npu': | ||
| return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)['num_vectorcore'] | ||
| else: | ||
| return 1 |
There was a problem hiding this comment.
This implementation of get_multiprocessor_count is incorrect for non-NPU devices. It returns 1 for all other backends, including CUDA. This will cause the grid size for Triton kernels to be fixed at 1, preventing parallel execution across multiprocessors and leading to a severe performance degradation.
Additionally, adding this function here redefines the existing get_multiprocessor_count at line 416. You should merge the NPU-specific logic into a single, correct function.
Here is a suggested implementation that correctly handles both NPU and other backends:
| def get_multiprocessor_count(tensor_idx: int = 0) -> int: | |
| if triton.runtime.driver.active.get_current_target().backend == 'npu': | |
| return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)['num_vectorcore'] | |
| else: | |
| return 1 | |
| def get_multiprocessor_count(tensor_idx: int = 0) -> int: | |
| try: | |
| if triton.runtime.driver.active.get_current_target().backend == 'npu': | |
| return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)['num_vectorcore'] | |
| return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)['multiprocessor_count'] | |
| except Exception: | |
| return 1 |
| @@ -0,0 +1,1244 @@ | |||
| # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang | |||
There was a problem hiding this comment.
This file is almost an exact copy of fla/modules/fused_norm_gate.py. Duplicating large files for minor device-specific changes makes the code harder to maintain. Any bug fix or feature enhancement would need to be applied in two places.
The only significant difference seems to be the autotune configuration for layer_norm_gated_fwd_kernel. This can be handled conditionally within the original fused_norm_gate.py file.
I recommend removing this file and modifying fla/modules/fused_norm_gate.py to support NPU-specific configurations. You can use fla.utils.device_platform to check the device and adjust the autotune parameters accordingly.
For example:
from fla.utils import device_platform
if device_platform == 'npu':
BT_VALUES = [32, 64]
else:
BT_VALUES = [16, 32, 64]
@triton.autotune(
configs=[
triton.Config({'BT': BT}, num_warps=num_warps)
for BT in BT_VALUES
for num_warps in [4, 8, 16]
],
...
)
def layer_norm_gated_fwd_kernel(...):
...This will make the codebase much cleaner and easier to maintain.
zhiyuan1i
force-pushed
the
main
branch
2 times, most recently
from
2b3db51 to
53dda79
Compare
1 participant
[WIP]Add NPU support for the
fused_norm_gateoperator