Skip to content

Wmma support for gemm_bias_add_reduce - #3316

Merged
illsilin merged 11 commits into
developfrom
streamhpc/gemm_bias_add_reduce_wmma
Jan 7, 2026
Merged

Wmma support for gemm_bias_add_reduce#3316
illsilin merged 11 commits into
developfrom
streamhpc/gemm_bias_add_reduce_wmma

Conversation

@EnricoDeg

Copy link
Copy Markdown
Contributor

Proposed changes

Summary:

  • Change EpilogueReduceCShuffle to support bias + add operations before reduction (multiple Ds)
  • Add wmma device struct for gemm_bias_add_reduce
  • Add instances (xdl parity)
  • Add tests

Checklist

Please put an x into the boxes that apply. You can also fill these out after creating the PR. If you're not sure, please don't hesitate to ask.

  • I have added tests relevant to the introduced functionality, and the unit tests are passing locally
  • I have added the test to REGRESSION_TESTS list defined at the top of CMakeLists.txt in tests/CMakeLists.txt, IF the test takes more than 30 seconds to run.
  • I have added inline documentation which enables the maintainers with understanding the motivation
  • I have removed the stale documentation which is no longer relevant after this pull request
  • (If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request
  • I have run clang-format on all changed files
  • Any dependent changes have been merged

Discussion

If this is a relatively large or complex change, feel free to start a discussion by explaining why you chose the solution you did and what alternatives you considered

Comment thread profiler/include/profiler/profile_gemm_bias_add_reduce_impl.hpp Outdated
@krithalith

Copy link
Copy Markdown
Contributor

In the profiler impl we have:

std::size_t num_byte = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N +
                       sizeof(CDataType) * M * N + sizeof(BiasDataType) * M * N +
                       sizeof(D0DataType) * M * N + sizeof(ReduceDataType) * M +
                       sizeof(ReduceDataType) * M;

But I thought the Bias was a simple 1D vector of size N?

Comment thread include/ck/tensor_operation/gpu/grid/epilogue_cshuffle_v3_reduce_wmma.hpp Outdated
Comment thread test/gemm_bias_add_reduce/test_gemm_common.hpp Outdated
@EnricoDeg
EnricoDeg force-pushed the streamhpc/gemm_bias_add_reduce_wmma branch 2 times, most recently from 7cd0696 to d7cbde3 Compare December 9, 2025 08:34
@EnricoDeg

Copy link
Copy Markdown
Contributor Author

In the profiler impl we have:

std::size_t num_byte = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N +
                       sizeof(CDataType) * M * N + sizeof(BiasDataType) * M * N +
                       sizeof(D0DataType) * M * N + sizeof(ReduceDataType) * M +
                       sizeof(ReduceDataType) * M;

But I thought the Bias was a simple 1D vector of size N?

Yeah it's wrong. I've fixed it

@EnricoDeg
EnricoDeg force-pushed the streamhpc/gemm_bias_add_reduce_wmma branch 2 times, most recently from 5693d1b to 0c73ab7 Compare December 12, 2025 12:25
@krithalith
krithalith self-requested a review December 12, 2025 16:57
krithalith
krithalith previously approved these changes Dec 12, 2025
@EnricoDeg
EnricoDeg force-pushed the streamhpc/gemm_bias_add_reduce_wmma branch 4 times, most recently from 1ef174f to 7eafbda Compare December 23, 2025 09:31
@EnricoDeg
EnricoDeg force-pushed the streamhpc/gemm_bias_add_reduce_wmma branch from 7eafbda to ab649a9 Compare January 6, 2026 11:13
@illsilin
illsilin merged commit aad4cf0 into develop Jan 7, 2026
21 checks passed
@illsilin
illsilin deleted the streamhpc/gemm_bias_add_reduce_wmma branch January 7, 2026 18:27
shumway pushed a commit that referenced this pull request Jan 15, 2026
* Add tests for gemm_bias_add_reduce

* Initial working implementation

* Generalize implementation of reduce epilogue

* Add tests for all layouts

* Add instances

* Fix test archs

* Fix xdl bug

* Remove library/profiler duplications

* Fix num_byted error profiler

* Fix typos

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants