Skip to content

[CUDA] gather_mm#3414

Merged
zcbenz merged 1 commit intoml-explore:mainfrom
zcbenz:cuda-gather-mm
Apr 17, 2026
Merged

[CUDA] gather_mm#3414
zcbenz merged 1 commit intoml-explore:mainfrom
zcbenz:cuda-gather-mm

Conversation

@zcbenz
Copy link
Copy Markdown
Collaborator

@zcbenz zcbenz commented Apr 15, 2026

Implement gather_mm by fusing CUTLASS GEMM with a custom prologue.

Note that this implementation is completely unoptimized (e.g. no tensor core, no vectorized read) and only performs 0.3x of cuBLAS, we would need to write our own heuristics to dispatch to proper main loops and alignments, which I'm postponing a bit until we have more CUTLASS kernels that I can make a general design.

@zcbenz zcbenz changed the title [CUDA] gather_matmul [CUDA] gather_mm Apr 17, 2026
Copy link
Copy Markdown
Member

@angeloskath angeloskath left a comment

Choose a reason for hiding this comment

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

Looks great. Should be fairly easy to extend to at least Ampere tensor cores.

@zcbenz zcbenz merged commit d142de6 into ml-explore:main Apr 17, 2026
16 checks passed
@zcbenz zcbenz deleted the cuda-gather-mm branch April 17, 2026 07:53
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