Skip to content

Thin N GEMM kernels - #4640

Merged
zcbenz merged 2 commits into
ml-explore:mainfrom
Isalia20:thin-nax-gemm
Oct 8, 2026
Merged

zcbenz merged 2 commits into
ml-explore:mainfrom
Isalia20:thin-nax-gemm

Conversation

@Isalia20

@Isalia20 Isalia20 commented Oct 7, 2026

Copy link
Copy Markdown
Contributor

GEMMs with small count of N and large M are bound by reading A. On main, the kernel uses 64x128 NAX tiles which waste most of the columns. For the kernel added in this PR, each simdgroup computes a 32x32 block of output matrix.

A matrix tile stays in registers and B matrix is read from device memory with mpp. K is split across simdgroups per output block, if M is <2048 then K is split across 4 simdgroups in quarters of K. If M is >=2048 then K is split in 2 main parts with 4 simdgroups while tile size used to write to output being bigger, so 4 simdgroups cover 64 columns instead of 32 columns

bf16 perf on M5 Pro with macOS 27:

Model, layer M x N x K main branch speedup
Qwen3.8-27B in_proj_a/b 2048 x 48 x 5120 106.1 65.1 1.63x
GLM-5.3 indexer weights_proj 2048 x 32 x 6144 122.6 65.1 1.88x
gpt-oss-20b router, bias 2048 x 32 x 2880 48.0 33.2 1.45x
Wan2.2-T2V-A14B proj_out, bias 32760 x 64 x 5120 1715.7 1324.9 1.29x

I don't have base m5 so would be nice to benchmark for it as well (M5 base is also included in the dispatch for this kernel).

  • AI usage disclosure:

AI was used to port kernels from my repo https://github.com/Isalia20/metalBLAS to here.

@zcbenz zcbenz left a comment

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 is a really nice optimization, thanks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants