enable llama3.1-8B on xpu #9434
Conversation
huaiyuzh
requested review from
merrymercy,
Ying1123,
hnyls2002,
xiezhq-hermann,
zhyncs,
ispobock,
HaiShaw,
ch-wan,
BBuf,
kushanam and
Edwardf0t1
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @huaiyuzh, 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 introduces support for the meta-llama/Llama-3.1-8B-Instruct model with BF16 datatype on XPU platforms. The changes primarily involve integrating XPU detection, implementing XPU-specific kernel dispatches for custom operations and activation functions, and updating system-wide checks to ensure compatibility and optimal performance on XPU hardware. This enhancement broadens the hardware compatibility for the specified Llama model, enabling its use on Intel's XPU architecture.
Highlights
- XPU Device Integration: Introduced detection for XPU devices and integrated XPU-specific dispatching logic for custom operations, allowing the system to leverage XPU hardware for optimized performance.
- XPU-Optimized Activation Functions: Added XPU-specific implementations for key activation functions like SiLU and GELU, utilizing
sgl_kernelfor efficient computation on XPU platforms. - Expanded
sgl-kernelSupport: Extended the availability checks forsgl-kernelimplementations across activation and layernorm layers to explicitly include XPU, ensuring proper kernel selection and fallback mechanisms. - Memory Pool Host Adjustments: Adjusted the memory pool host logic to correctly account for XPU, modifying conditional imports for
sgl_kernel.kvcacheiofunctions to align with XPU compatibility.
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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| 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 issue 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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request adds support for XPU devices by introducing is_xpu checks and forward_xpu implementations in several layers. The changes are generally well-structured. My review includes a critical fix for a potential bug in SiluAndMul's XPU forward pass, along with several suggestions to improve code maintainability by reducing duplication and enhancing readability.
| d = x.shape[-1] // 2 | ||
| output_shape = x.shape[:-1] + (d,) | ||
| out = torch.empty(output_shape, dtype=x.dtype, device=x.device) | ||
| out = silu_and_mul(x, out) |
There was a problem hiding this comment.
The line out = silu_and_mul(x, out) is inconsistent with the forward_cuda implementation and may cause a bug. If silu_and_mul modifies out in-place and returns None (which is a common pattern for such functions for performance), this function will incorrectly return None.
To fix this, you should call silu_and_mul without assigning its return value, just like in forward_cuda.
| out = silu_and_mul(x, out) | |
| silu_and_mul(x, out) |
| if _is_xpu: | ||
| from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul |
There was a problem hiding this comment.
This block duplicates the imports from the _is_cuda case. For better maintainability, you could consider combining these conditions in the future, for example: if _is_cuda or _is_xpu:.
| def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: | ||
| d = x.shape[-1] // 2 | ||
| output_shape = x.shape[:-1] + (d,) | ||
| out = torch.empty(output_shape, dtype=x.dtype, device=x.device) | ||
| if self.approximate == "tanh": | ||
| gelu_tanh_and_mul(x, out) | ||
| elif self.approximate == "none": | ||
| gelu_and_mul(x, out) | ||
| else: | ||
| raise RuntimeError("GeluAndMul only support tanh or none") | ||
| return out |
There was a problem hiding this comment.
The implementation of forward_xpu is identical to forward_cuda. To improve maintainability and reduce code duplication, consider creating a shared helper method that both forward_xpu and forward_cuda can call.
For example, you could create a _forward_impl method and have both forward_cuda and forward_xpu delegate to it. This would make the code cleaner and easier to maintain in the future.
|
|
||
| _is_npu = is_npu() | ||
| if not _is_npu: | ||
| _is_xpu = is_xpu() | ||
| if (not _is_npu and not _is_xpu): |
There was a problem hiding this comment.
The condition (not _is_npu and not _is_xpu) can be simplified using De Morgan's laws to not (_is_npu or _is_xpu) for better readability. The outer parentheses are also redundant.
| if (not _is_npu and not _is_xpu): | |
| if not (_is_npu or _is_xpu): |
This PR is to enable meta-llama/Llama-3.1-8B-Instruct with datatype BF16 on XPU.