feat(kernel): Add batch matmul API#668
Conversation
qedawkins
changed the title
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ceb451dc62
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
ceb451d to
0d34978
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0d349780f3
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
0148a86 to
ba250a5
Compare
|
@qedawkins the gluon bf16 gemm registration has been removed in #689 due to performance issue in decode. can we also follow it to make sure not registering it until performance is fixed? just for normal mm not batch mm |
ba250a5 to
99d1eb0
Compare
I kept bmm disabled too for consistency. The situations where we're faster are quite narrow for bmm also so fine to enable slowly there as well. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 99d1eb057c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 459b46a9a5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 99327b138e
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
cb3e56a to
d6d5ca3
Compare
Add a public gemm.bmm API matching the GEMM dispatch style for true batched inputs. The API is for true batched GEMMs only and is kept distinct from typical GEMM as implementations can vary in non-trivial ways depending on the size of the batch. Currently registers the torch reference implementation while performant backend implementations are left for follow-up work. Also add out= plumbing to mm and bmm with torch reference dense GEMM/BMM write through native torch out= paths when possible. Triton GEMM paths that already store through explicit output strides write directly into out=. Gluon GEMM paths were updated to do the same as well. This is motivated by bmm callsites needing to write directly into concatenated outputs in ways that wasn't yet encountered for normal GEMM. Bias and enable_pdl handling now follow the same ownership model for mm and bmm: kernels that opt in receive those kwargs, and dispatch only applies bias for kernels without bias support. Share the GEMM and BMM numerics input generation paths so shape/layout handling stays consistent, and add shape presets, throughput accounting, API selection coverage, generator coverage, and focused API correctness tests. Testing: pytest tokenspeed-kernel/test/ops/test_bmm.py tokenspeed-kernel/test/ops/test_gemm.py; pytest tokenspeed-kernel/test/test_kernel_api_selection.py; pytest tokenspeed-kernel/test/test_numerics.py; pytest tokenspeed-kernel/test/test_benchmark.py -q -k 'bmm_throughput'. Signed-off-by: Quinn Dawkins <quinn.dawkins@gmail.com>
d6d5ca3 to
cc76cfb
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cc76cfb967
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Signed-off-by: Quinn Dawkins <quinn.dawkins@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 801ad47a7c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
3 participants
Add a public gemm.bmm API matching the GEMM dispatch style for true batched inputs. The API is for true batched GEMMs only and is kept distinct from typical GEMM as implementations can vary in non-trivial ways depending on the size of the batch. Currently registers the torch reference implementation while performant backend implementations are left for follow-up work.
Also add out= plumbing to mm and bmm with torch reference dense GEMM/BMM write through native torch out= paths when possible. Triton GEMM paths that already store through explicit output strides write directly into out=. Gluon GEMM paths were updated to do the same as well. This is motivated by bmm callsites needing to write directly into concatenated outputs in ways that wasn't yet encountered for normal GEMM. Bias and enable_pdl handling now follow the same ownership model for mm and bmm: kernels that opt in receive those kwargs, and dispatch only applies bias for kernels without bias support.
Share the GEMM and BMM numerics input generation paths so shape/layout handling stays consistent, and add shape presets, throughput accounting, API selection coverage, generator coverage, and focused API correctness tests.
Testing: pytest tokenspeed-kernel/test/ops/test_bmm.py tokenspeed-kernel/test/ops/test_gemm.py; pytest tokenspeed-kernel/test/test_kernel_api_selection.py; pytest tokenspeed-kernel/test/test_numerics.py; pytest tokenspeed-kernel/test/test_benchmark.py -q -k 'bmm_throughput'.