Skip to content

Wmma support for gemm_blockscale_wp - #3315

Merged
illsilin merged 2 commits into
streamhpc/gemm_ab_scalefrom
streamhpc/gemm_blockscale_wp_wmma
Dec 3, 2025
Merged

Wmma support for gemm_blockscale_wp#3315
illsilin merged 2 commits into
streamhpc/gemm_ab_scalefrom
streamhpc/gemm_blockscale_wp_wmma

Conversation

@EnricoDeg

@EnricoDeg EnricoDeg commented Nov 27, 2025

Copy link
Copy Markdown
Contributor

Proposed changes

Summary:

  • device: wmma blockscale preshuffle implementation
  • device: add support for preshuffle kernels dispatch in common invoker to avoid code duplications
  • block: add pipeline v1 with ab scale and LDS skip for B matrix (it should be possible to unify it with the existing pipeline with LDS skip for B, but it will be target in another issue once both pipelines have been optimized)
  • grid: add preshuffle support in GridwiseGemm_wmma_cshuffle_v3_ab_scale
  • grid: add AScaleLayout and BScaleLayout template parameters to support different layouts between A matrix and A scale matrix (same for B)
  • add example and instances

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

@EnricoDeg
EnricoDeg changed the base branch from develop to streamhpc/gemm_ab_scale November 27, 2025 08:31
@EnricoDeg
EnricoDeg force-pushed the streamhpc/gemm_blockscale_wp_wmma branch 2 times, most recently from 1dbf2b8 to 28f1135 Compare November 27, 2025 10:04
@EnricoDeg EnricoDeg self-assigned this Nov 27, 2025
@EnricoDeg
EnricoDeg force-pushed the streamhpc/gemm_blockscale_wp_wmma branch from 4ad5bd7 to d15a0dc Compare November 28, 2025 12:28
@EnricoDeg
EnricoDeg force-pushed the streamhpc/gemm_ab_scale branch from 97badf1 to 647f08c Compare December 1, 2025 17:48
@EnricoDeg
EnricoDeg force-pushed the streamhpc/gemm_blockscale_wp_wmma branch from 89e6349 to 2793d3e Compare December 1, 2025 17:53
@EnricoDeg
EnricoDeg force-pushed the streamhpc/gemm_ab_scale branch from 647f08c to a023f2c Compare December 2, 2025 09:23
@EnricoDeg
Support for preshuffle with ab scale
d5898d9 
 - add support for b preshuffle in GridwiseGemm_wmma_cshuffle_v3_ab_scale
 - add support for AScaleLayout amnd BScaleLayout (can be different
   from ALayout and BLayout, respectively)
 - add Run method in v1 pipeline to support preshuffle + scaling
 - add support for preshuffle gemms in common invoker
 - Add splitk support
@EnricoDeg
EnricoDeg force-pushed the streamhpc/gemm_blockscale_wp_wmma branch from 2793d3e to d5898d9 Compare December 2, 2025 09:39
@illsilin
illsilin merged commit c8ec0d0 into streamhpc/gemm_ab_scale Dec 3, 2025
17 checks passed
@illsilin
illsilin deleted the streamhpc/gemm_blockscale_wp_wmma branch December 3, 2025 16:04
EnricoDeg added a commit that referenced this pull request Dec 3, 2025
@EnricoDeg
Wmma support for gemm_blockscale_wp (#3315)
8e62dc9 
* Support for  preshuffle with ab scale

 - add support for b preshuffle in GridwiseGemm_wmma_cshuffle_v3_ab_scale
 - add support for AScaleLayout amnd BScaleLayout (can be different
   from ALayout and BLayout, respectively)
 - add Run method in v1 pipeline to support preshuffle + scaling
 - add support for preshuffle gemms in common invoker
 - Add splitk support

* Fix copyright header
EnricoDeg added a commit that referenced this pull request Dec 9, 2025
@EnricoDeg
Wmma support for gemm_blockscale_wp (#3315)
a29cf5d 
* Support for  preshuffle with ab scale

 - add support for b preshuffle in GridwiseGemm_wmma_cshuffle_v3_ab_scale
 - add support for AScaleLayout amnd BScaleLayout (can be different
   from ALayout and BLayout, respectively)
 - add Run method in v1 pipeline to support preshuffle + scaling
 - add support for preshuffle gemms in common invoker
 - Add splitk support

* Fix copyright header
EnricoDeg added a commit that referenced this pull request Dec 9, 2025
@EnricoDeg
Wmma support for gemm_blockscale_wp (#3315)
a0e762e 
* Support for  preshuffle with ab scale

 - add support for b preshuffle in GridwiseGemm_wmma_cshuffle_v3_ab_scale
 - add support for AScaleLayout amnd BScaleLayout (can be different
   from ALayout and BLayout, respectively)
 - add Run method in v1 pipeline to support preshuffle + scaling
 - add support for preshuffle gemms in common invoker
 - Add splitk support

* Fix copyright header
EnricoDeg added a commit that referenced this pull request Dec 11, 2025
@EnricoDeg
Wmma support for gemm_ab_scale (#3314)
ce99cab 
* Support gemm_ab_scale:

 - Add tests
 - Integrate scaling implementation in multiple D
 - Generalize existing b_scale for ab_scale
 - Add instances
 - Generalize implementation for ScaleBlockM, ScaleBlockN, ScaleBlockK
 - Add support for all layouts supported by xdl
 - Fix splitk xdl

* Fix copyright

* Wmma support for gemm_blockscale_wp (#3315)

* Support for  preshuffle with ab scale

 - add support for b preshuffle in GridwiseGemm_wmma_cshuffle_v3_ab_scale
 - add support for AScaleLayout amnd BScaleLayout (can be different
   from ALayout and BLayout, respectively)
 - add Run method in v1 pipeline to support preshuffle + scaling
 - add support for preshuffle gemms in common invoker
 - Add splitk support

* Fix copyright header
kabrahamAMD pushed a commit that referenced this pull request Dec 11, 2025
@EnricoDeg @kabraham-streamhpc
Wmma support for gemm_ab_scale (#3314)
c37e4bf 
* Support gemm_ab_scale:

 - Add tests
 - Integrate scaling implementation in multiple D
 - Generalize existing b_scale for ab_scale
 - Add instances
 - Generalize implementation for ScaleBlockM, ScaleBlockN, ScaleBlockK
 - Add support for all layouts supported by xdl
 - Fix splitk xdl

* Fix copyright

* Wmma support for gemm_blockscale_wp (#3315)

* Support for  preshuffle with ab scale

 - add support for b preshuffle in GridwiseGemm_wmma_cshuffle_v3_ab_scale
 - add support for AScaleLayout amnd BScaleLayout (can be different
   from ALayout and BLayout, respectively)
 - add Run method in v1 pipeline to support preshuffle + scaling
 - add support for preshuffle gemms in common invoker
 - Add splitk support

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

Reviewers

@illsilin illsilin illsilin approved these changes
@carlushuang carlushuang Awaiting requested review from carlushuang
@qianfengz qianfengz Awaiting requested review from qianfengz
@aosewski aosewski Awaiting requested review from aosewski
@poyenc poyenc Awaiting requested review from poyenc
@geyyer geyyer Awaiting requested review from geyyer
@bartekxk bartekxk Awaiting requested review from bartekxk
@andriy-ca andriy-ca Awaiting requested review from andriy-ca
@afagaj afagaj Awaiting requested review from afagaj
@asleepzzz asleepzzz Awaiting requested review from asleepzzz
@ThomasNing ThomasNing Awaiting requested review from ThomasNing
@coderfeli coderfeli Awaiting requested review from coderfeli
@cgmillette cgmillette Awaiting requested review from cgmillette
@shumway shumway Awaiting requested review from shumway
@vidyasagar-amd vidyasagar-amd Awaiting requested review from vidyasagar-amd
@tenpercent tenpercent Awaiting requested review from tenpercent

Assignees

@EnricoDeg EnricoDeg

Labels

organization: streamhpc

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@EnricoDeg @illsilin