Implement swizzle gemm for flux dit and wan dit #204
Conversation
Summary of ChangesHello @guangzlu, 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 performance of linear operations within the Highlights
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
|
Summary of ChangesHello @guangzlu, 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 a significant performance enhancement for the Highlights
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 introduces a swizzle gemm optimization for the flux dit and wan dit models. The core logic is encapsulated in a new file, diffsynth_engine/utils/aiter_linear.py, which provides a context manager to monkey-patch torch.nn.functional.linear. The changes in the model files simply enable this new functionality. My review focuses on the new implementation, where I've identified a critical issue with a hardcoded device that could cause problems in multi-GPU environments. I have provided suggestions to fix this and to refactor the code for better maintainability.
| def optimized_linear(input, weight, bias=None, otype=torch.bfloat16, | ||
| scaleA=None, scaleB=None, device="cuda"): | ||
|
|
||
| input_flat = input.reshape(-1, input.shape[-1]) | ||
|
|
||
| init_hipblas() | ||
|
|
||
| weight_preshuffle = shuffle_weight(weight.contiguous(), layout=(16, 16), use_int4=False).to(device) | ||
| output_flat = hipb_mm( | ||
| input_flat, | ||
| weight_preshuffle.t(), | ||
| bias=bias, | ||
| solution_index=-1, | ||
| out_dtype=otype, | ||
| scaleA=scaleA, | ||
| scaleB=scaleB, | ||
| scaleOut=None, | ||
| bpreshuffle=True | ||
| ) | ||
|
|
||
| # Reshape output to match input dimensions | ||
| new_shape = input.shape[:-1] + (weight.shape[0],) | ||
| output = output_flat.view(new_shape) | ||
| return output |
There was a problem hiding this comment.
The device parameter is hardcoded to "cuda", which will cause errors if the model is run on a different device (e.g., cuda:1). The device should be dynamically determined from the input tensor. The function signature should also be updated to not include the device parameter, to better align with the original torch.nn.functional.linear.
def optimized_linear(input, weight, bias=None, otype=torch.bfloat16,
scaleA=None, scaleB=None):
input_flat = input.reshape(-1, input.shape[-1])
init_hipblas()
weight_preshuffle = shuffle_weight(weight.contiguous(), layout=(16, 16), use_int4=False).to(input.device)
output_flat = hipb_mm(
input_flat,
weight_preshuffle.t(),
bias=bias,
solution_index=-1,
out_dtype=otype,
scaleA=scaleA,
scaleB=scaleB,
scaleOut=None,
bpreshuffle=True
)
# Reshape output to match input dimensions
new_shape = input.shape[:-1] + (weight.shape[0],)
output = output_flat.view(new_shape)
return output| def optimized_linear_fp8(input, weight, bias=None, otype=torch.bfloat16, | ||
| scaleA=None, scaleB=None, device="cuda"): | ||
|
|
||
| input_flat = input.reshape(-1, input.shape[-1]) | ||
|
|
||
| if use_scale_for_fp8: | ||
|
|
||
| input_flat, a_scale = per_tensor_quant_hip(input_flat, quant_dtype=DTYPE_FP8) | ||
| weight = weight.to(DTYPE_FP8) | ||
|
|
||
| init_hipblas() | ||
|
|
||
| weight_preshuffle = shuffle_weight(weight.contiguous(), layout=(16, 16)).to(device) | ||
| output_flat = hipb_mm( | ||
| input_flat, | ||
| weight_preshuffle.t(), | ||
| bias=bias, | ||
| solution_index=-1, | ||
| out_dtype=otype, | ||
| scaleA=a_scale, | ||
| scaleB=scaleB, | ||
| scaleOut=None, | ||
| bpreshuffle=True | ||
| ) | ||
|
|
||
| else: | ||
| input_flat = input_flat.to(DTYPE_FP8) | ||
| weight = weight.to(DTYPE_FP8) | ||
|
|
||
| init_hipblas() | ||
|
|
||
| weight_preshuffle = shuffle_weight(weight.contiguous(), layout=(16, 16)).to(device) | ||
| output_flat = hipb_mm( | ||
| input_flat, | ||
| weight_preshuffle.t(), | ||
| bias=bias, | ||
| solution_index=-1, | ||
| out_dtype=otype, | ||
| scaleA=scaleA, | ||
| scaleB=scaleB, | ||
| scaleOut=None, | ||
| bpreshuffle=True | ||
| ) | ||
|
|
||
|
|
||
| # Reshape output to match input dimensions | ||
| new_shape = input.shape[:-1] + (weight.shape[0],) | ||
| output = output_flat.view(new_shape) | ||
| return output |
There was a problem hiding this comment.
This function has a couple of issues:
- Similar to
optimized_linear, it hardcodesdevice="cuda", which can lead to runtime errors on multi-GPU setups. The device should be inferred from the input tensor. - There is significant code duplication between the
ifandelsebranches. This can be refactored to improve readability and maintainability.
def optimized_linear_fp8(input, weight, bias=None, otype=torch.bfloat16,
scaleA=None, scaleB=None):
input_flat = input.reshape(-1, input.shape[-1])
device = input.device
if use_scale_for_fp8:
input_flat, a_scale = per_tensor_quant_hip(input_flat, quant_dtype=DTYPE_FP8)
scaleA = a_scale
else:
input_flat = input_flat.to(DTYPE_FP8)
weight = weight.to(DTYPE_FP8)
init_hipblas()
weight_preshuffle = shuffle_weight(weight.contiguous(), layout=(16, 16)).to(device)
output_flat = hipb_mm(
input_flat,
weight_preshuffle.t(),
bias=bias,
solution_index=-1,
out_dtype=otype,
scaleA=scaleA,
scaleB=scaleB,
scaleOut=None,
bpreshuffle=True
)
# Reshape output to match input dimensions
new_shape = input.shape[:-1] + (weight.shape[0],)
output = output_flat.view(new_shape)
return output
There was a problem hiding this comment.
Code Review
This pull request introduces swizzle GEMM optimizations for the flux dit and wan dit models, aiming to enhance performance, particularly on AMD hardware. This is accomplished by monkey-patching torch.nn.functional.linear within a new context manager. The implementation is sound, but I have identified a few areas for improvement to increase code quality, robustness, and configurability. My review includes suggestions to remove an unused import, refactor duplicated code, and address a hardcoded device name in aiter_linear.py which could lead to issues in multi-GPU setups. Additionally, I've recommended making the swizzle optimization a configurable option in flux_dit.py and wan_dit.py rather than being hardcoded.
| def optimized_linear(input, weight, bias=None, otype=torch.bfloat16, | ||
| scaleA=None, scaleB=None, device="cuda"): | ||
|
|
||
| input_flat = input.reshape(-1, input.shape[-1]) | ||
|
|
||
| init_hipblas() | ||
|
|
||
| weight_preshuffle = shuffle_weight(weight.contiguous(), layout=(16, 16), use_int4=False).to(device) |
There was a problem hiding this comment.
The device parameter defaults to a hardcoded "cuda". This can cause runtime errors if the model is deployed on a different device (e.g., cuda:1). It's more robust to determine the device from the input tensor. Additionally, adding a non-standard device parameter to a function that monkey-patches F.linear can be risky. I suggest removing it and using input.device instead.
| def optimized_linear(input, weight, bias=None, otype=torch.bfloat16, | |
| scaleA=None, scaleB=None, device="cuda"): | |
| input_flat = input.reshape(-1, input.shape[-1]) | |
| init_hipblas() | |
| weight_preshuffle = shuffle_weight(weight.contiguous(), layout=(16, 16), use_int4=False).to(device) | |
| def optimized_linear(input, weight, bias=None, otype=torch.bfloat16, | |
| scaleA=None, scaleB=None): | |
| input_flat = input.reshape(-1, input.shape[-1]) | |
| init_hipblas() | |
| weight_preshuffle = shuffle_weight(weight.contiguous(), layout=(16, 16), use_int4=False).to(input.device) |
| def optimized_linear_fp8(input, weight, bias=None, otype=torch.bfloat16, | ||
| scaleA=None, scaleB=None, device="cuda"): | ||
|
|
||
| input_flat = input.reshape(-1, input.shape[-1]) | ||
|
|
||
| if use_scale_for_fp8: | ||
|
|
||
| input_flat, a_scale = per_tensor_quant_hip(input_flat, quant_dtype=DTYPE_FP8) | ||
| weight = weight.to(DTYPE_FP8) | ||
|
|
||
| init_hipblas() | ||
|
|
||
| weight_preshuffle = shuffle_weight(weight.contiguous(), layout=(16, 16)).to(device) | ||
| output_flat = hipb_mm( | ||
| input_flat, | ||
| weight_preshuffle.t(), | ||
| bias=bias, | ||
| solution_index=-1, | ||
| out_dtype=otype, | ||
| scaleA=a_scale, | ||
| scaleB=scaleB, | ||
| scaleOut=None, | ||
| bpreshuffle=True | ||
| ) | ||
|
|
||
| else: | ||
| input_flat = input_flat.to(DTYPE_FP8) | ||
| weight = weight.to(DTYPE_FP8) | ||
|
|
||
| init_hipblas() | ||
|
|
||
| weight_preshuffle = shuffle_weight(weight.contiguous(), layout=(16, 16)).to(device) | ||
| output_flat = hipb_mm( | ||
| input_flat, | ||
| weight_preshuffle.t(), | ||
| bias=bias, | ||
| solution_index=-1, | ||
| out_dtype=otype, | ||
| scaleA=scaleA, | ||
| scaleB=scaleB, | ||
| scaleOut=None, | ||
| bpreshuffle=True | ||
| ) |
There was a problem hiding this comment.
This function can be improved in two ways:
- Hardcoded Device: Similar to
optimized_linear, thedeviceparameter is hardcoded to"cuda", which can lead to errors in a multi-GPU environment. It's better to derive the device frominput.device. - Code Duplication: There's significant code duplication between the
if use_scale_for_fp8:andelsebranches. Refactoring this will improve readability and maintainability.
I've provided a suggestion that addresses both points.
def optimized_linear_fp8(input, weight, bias=None, otype=torch.bfloat16,
scaleA=None, scaleB=None):
input_flat = input.reshape(-1, input.shape[-1])
a_scale_for_mm = scaleA
if use_scale_for_fp8:
input_flat, a_scale_for_mm = per_tensor_quant_hip(input_flat, quant_dtype=DTYPE_FP8)
else:
input_flat = input_flat.to(DTYPE_FP8)
weight = weight.to(DTYPE_FP8)
init_hipblas()
weight_preshuffle = shuffle_weight(weight.contiguous(), layout=(16, 16)).to(input.device)
output_flat = hipb_mm(
input_flat,
weight_preshuffle.t(),
bias=bias,
solution_index=-1,
out_dtype=otype,
scaleA=a_scale_for_mm,
scaleB=scaleB,
scaleOut=None,
bpreshuffle=True
)| use_cfg = hidden_states.shape[0] > 1 | ||
| with ( | ||
| fp8_inference(fp8_linear_enabled), | ||
| use_swizzle_hipblaslt(swizzle=True, use_fp8_linear=fp8_linear_enabled), |
There was a problem hiding this comment.
The swizzle parameter is hardcoded to True. This makes it difficult to disable this optimization without changing the code. Consider making this configurable, for example, through a class attribute self.swizzle_enabled that can be set during model initialization.
| use_cfg = x.shape[0] > 1 | ||
| with ( | ||
| fp8_inference(fp8_linear_enabled), | ||
| use_swizzle_hipblaslt(swizzle=True, use_fp8_linear=fp8_linear_enabled), |
There was a problem hiding this comment.
The swizzle parameter is hardcoded to True. This makes it difficult to disable this optimization without changing the code. Consider making this configurable, for example, through a class attribute self.swizzle_enabled that can be set during model initialization.
| import torch.nn.functional as F | ||
| from functools import lru_cache | ||
| from aiter import hipb_mm, hipb_create_extension, per_tensor_quant_hip | ||
| from aiter.tuned_gemm import tgemm |
There was a problem hiding this comment.
This import of tgemm is not used within the file and can be safely removed to improve code cleanliness.
1 participant
Implement swizzle gemm for flux dit and wan dit.