Skip to content

metal: threadgroup-memory alignment and expert-index width fixes #27944

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

Closed
tarruda wants to merge 2 commits into from
Closed

metal: threadgroup-memory alignment and expert-index width fixes #27944

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b0af918
metal: align dynamic threadgroup memory
tarruda Aug 28, 2026
b5cb930
metal: widen expert matmul thread index
tarruda Aug 28, 2026
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
2 changes: 1 addition & 1 deletion ggml/src/ggml-metal/ggml-metal-device.m
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -800,7 +800,7 @@ void ggml_metal_encoder_set_buffer(ggml_metal_encoder_t encoder, struct ggml_met
}

void ggml_metal_encoder_set_threadgroup_memory_size(ggml_metal_encoder_t encoder, size_t size, int idx) {
[encoder->obj setThreadgroupMemoryLength:size atIndex:idx];
[encoder->obj setThreadgroupMemoryLength:GGML_PAD(size, 16) atIndex:idx];
}
Comment on lines 802 to 804

@ggerganov ggerganov Aug 29, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

I have extracted this change into: #27951


void ggml_metal_encoder_dispatch_threadgroups(ggml_metal_encoder_t encoder, int tg0, int tg1, int tg2, int tptg0, int tptg1, int tptg2) {
Expand Down
2 changes: 1 addition & 1 deletion ggml/src/ggml-metal/kernels/mul_mm.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -435,7 +435,7 @@ kernel void kernel_mul_mm_id(
device char * dst,
threadgroup char * shmem [[threadgroup(0)]],
uint3 tgpig[[threadgroup_position_in_grid]],
ushort tiitg[[thread_index_in_threadgroup]],
uint tiitg[[thread_index_in_threadgroup]],

@ggerganov ggerganov Aug 29, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This change does not seem necessary, so for now not merging it. That is unless the AI can demonstrate an issue with it that I am not seeing.

ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
threadgroup S0 * sa = (threadgroup S0 *)(shmem);
Expand Down
1 change: 1 addition & 0 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9359,6 +9359,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
// test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 512, 262144, 9216, {1, 1}, {1, 1}));

// test large experts*tokens
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 512, 10, false, 64, 512, 64));

@ggerganov ggerganov Aug 29, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

@tarruda Thanks. And are you sure this test fails on master? It passes on my end without the metal changes in the PR.

@tarruda tarruda Aug 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Replied here

Seems like it requires some validation to be enabled for it to be detected on master

for (bool b : {false, true}) {
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, b, 32, 1024, 16));
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 2, 2, b, 32, 8192, 64));
Expand Down