-
Notifications
You must be signed in to change notification settings - Fork 13
Restructure MoE and Add MoE prepare input kernels #29
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.
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
base: main
Are you sure you want to change the base?
Conversation
- restructure moe kernels folder - add prepare moe inputs kerel Signed-off-by: kareem <[email protected]>
Signed-off-by: kareem <[email protected]>
Signed-off-by: kareem <[email protected]>
Signed-off-by: kareem <[email protected]>
Signed-off-by: kareem <[email protected]>
Signed-off-by: kareem <[email protected]>
Signed-off-by: Shaik, Kareem M <[email protected]>
Signed-off-by: kareem <[email protected]>
Signed-off-by: kareem <[email protected]>
adityachatter
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
activations not only serve for MoE, better leave it unchanged.
Beside I'd prefer only put customized cutlass code under src/sycl/kernels/, and leave pure SYCL code outside
|
|
||
| @pytest.mark.parametrize("num_tokens", [5, 16, 128]) | ||
| @pytest.mark.parametrize("num_experts", [4, 8, 32]) | ||
| @pytest.mark.parametrize("top_k", [2]) | ||
| @pytest.mark.parametrize("hidden_dims", [16, 32, 64]) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
| void operator()(sycl::nd_item<1> item) const { | ||
| int32_t tot_offset = 0; | ||
| int32_t tot_rounded_offset = 0; | ||
| expert_offsets_[0] = 0; | ||
| blockscale_offsets_[0] = 0; | ||
| for (int i = 0; i < num_experts_; ++i) { | ||
| atomic_buffer_[i] = tot_offset; | ||
| int num_tokens = problem_sizes1_[i * 3]; | ||
| int rounded_num_tokens = (num_tokens + (block_size - 1)) / block_size * block_size; | ||
| tot_offset += num_tokens; | ||
| tot_rounded_offset += rounded_num_tokens; | ||
| expert_offsets_[i + 1] = tot_offset; | ||
| blockscale_offsets_[i + 1] = tot_rounded_offset; | ||
| } | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The function is purely sequential? Try to parallelize it using sycl::exclusive_scan, see https://github.com/intel/llvm/blob/4474e85c51c1c3153af9938164391d1e836cfff4/sycl/doc/extensions/removed/sycl_ext_oneapi_group_algorithms.asciidoc?plain=1#L75
Uh oh!
There was an error while loading. Please reload this page.