Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[GPU] Optimize RMS Stack Size for Better Performance #26515

Merged
merged 11 commits into from
Sep 11, 2024
Merged

[GPU] Optimize RMS Stack Size for Better Performance #26515

Changes from 1 commit
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
ace32d9
[GPU] optimize RMS kernel
zaixing-wang Sep 10, 2024
e3680ff
consider larger dataSize && leftover
zaixing-wang Sep 10, 2024
5af25e2
use const size_t to instead Macro
zaixing-wang Sep 10, 2024
ad8abda
Merge remote-tracking branch 'upstream/master' into rms_optim
zaixing-wang Sep 10, 2024
5cd1676
update comment
zaixing-wang Sep 10, 2024
1c5d0b4
fix codestyle
zaixing-wang Sep 10, 2024
40bbbd8
update get_item_num_and_lws
zaixing-wang Sep 11, 2024
9c7dbbc
clean
zaixing-wang Sep 11, 2024
216eef4
clean declare
zaixing-wang Sep 11, 2024
088d9f9
update
zaixing-wang Sep 11, 2024
ba6bb70
Merge remote-tracking branch 'upstream/master' into rms_optim
zaixing-wang Sep 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions src/plugins/intel_gpu/src/kernel_selector/kernels/rms/rms_kernel_bfyx_opt.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
#include "kernel_selector_utils.h"
#include <string>

#define MAX_ITEMS_NUM 8

namespace kernel_selector {
static constexpr size_t subgroup_size = 16;
ParamsKey RMSKernelBfyxOpt::GetSupportedKey() const {
Expand Down Expand Up @@ -58,9 +60,9 @@ JitConstants RMSKernelBfyxOpt::GetJitConstants(const rms_params& params, Dispatc
}

const std::string lws_0 = "get_local_size(0)";
// It can be expected that the maximum possible itemsNum will not exceed 32
// Therefore, in dynamic shape, stack_size including additional buffer is set to 33
constexpr size_t stack_size = 33;
// It can be expected that the maximum possible itemsNum will not exceed MAX_ITEMS_NUM
// Therefore, in dynamic shape, stack_size including additional buffer is set to MAX_ITEMS_NUM
constexpr size_t stack_size = MAX_ITEMS_NUM;
vladimir-paramuzov marked this conversation as resolved.
Show resolved Hide resolved
jit.AddConstants({
MakeJitConstant("DATA_SIZE", data_size),
MakeJitConstant("LWS", lws_0),
Expand Down Expand Up @@ -120,7 +122,7 @@ RMSKernelBase::DispatchData RMSKernelBfyxOpt::SetDefault(const rms_params& param

dispatchData.itemsNum = dispatchData.dataSize;
Copy link
Contributor

@dnkurek dnkurek Sep 10, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hint: maybe check if dispatchData.dataSize is actually constant and won't change when executing LLM, therefore we know beforehand what will always dispatchData.itemsNum be in each execution. In other words, maybe the only thing that changes is dataCount

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In Qwen, dataSize is always equal to 896.

Copy link
Contributor

@dnkurek dnkurek Sep 10, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, so that looks to me like a possible optimization there, since now you exactly know how much private memory you will need and better use the GPU's resources

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And also you would know LWS and therefore you could use reqd_work_group_size

Copy link
Contributor

@dnkurek dnkurek Sep 10, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also 896 seems relatively not a lot. It makes sense why reducing stack size would improve performance, since you are basically freeing up resources that are not used.

// Compute maximum possible LWS that does not exceed device capabilities and optimizes number of global memory reads
while ((dispatchData.itemsNum > 8 || dispatchData.lws[0] < dispatchData.itemsNum) && (2 * dispatchData.lws[0] <= max_lws)) {
while ((dispatchData.itemsNum > MAX_ITEMS_NUM || dispatchData.lws[0] < dispatchData.itemsNum) && (2 * dispatchData.lws[0] <= max_lws)) {
dispatchData.lws[0] *= 2;
dispatchData.itemsNum /= 2;
}
Expand Down
Loading