[CK_TILE] Extend support of mix precision microscaling BQuant by assistant-librarian[bot] · Pull Request #4267 · ROCm/rocm-libraries

assistant-librarian · 2026-02-03T21:46:04Z

Proposed changes

Supported types combinations using BQuant=e8m0:

A=bf16
B=bf16,bf8,fp4

Summary:

remove usage of pk_fp4_raw_t: consistent with other implementations and avoid taking into account of the packed size explicitly. In general, the raw type should not be used because CK Tile internally takes care of the PackedSize, so using the raw type adds unnecessary complexity to the implementation
handle microscaling by checking for e8m0 type for BQuant (previous implementation was inconsistent)
add support for scaling instructions in DequantPack8
mx pipeline:
- extend existing pipeline to support different B types
- add support to scale and cast before writing to LDS or after reading from LDS (this can be defined in the Problem by the user)
block gemm:
- mx pipeline is now using block gemm BQuant
- block gemm BQuant can now load from LDS and apply scale and then call block gemm universal operator. This adds new functionalities and remove code duplication
warp gemm:
- add case to support 128bit ds_read/write for both A and B when A=16bit and B=8bit
add examples and tests: note that some tests for bf16/fp4 already existed but were removed during previous tests refactoring. I added them again and other relevant tests for new types combinations

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

🔁 Imported from ROCm/composable_kernel#3689
🧑‍💻 Originally authored by @EnricoDeg

Both A and B are using 128 bit instructions

Support both scale before writing to LDS or scale after reading from LDS

…evelop/ROCm_composable_kernel/pr-3689

…3689

ThomasNing

Great PR! Thanks for the contribution Enrico!

…3689

## Motivation Extend functionalities of mix precision microscale BQuant. ## Technical Details - Support combinations of A 16 bit and B 16/8/4 bit using scale cvt instructions on gfx950 - Enable LDS load transpose for 4 bit types in ck tile - Allow to use LDS load transpose when A and B have different data types in LDS - Define `AttrNumAccessEnum` for A and B separately - If `AttrNumAccessA != AttrNumAccessB`, use packed NumAccess instead of interleaved one - Support all layouts combinations with `BCastAfterLDSRead` - Initial support for different layouts with `BCastBeforeLDSWrite`. BQuant encodings still need to be generalized (will be targeted in follow-up PR). - Add tests for different layouts and split them into multiple files Minor change to CK pre-commit config to fix the path (after transition to mono-repo) ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

…3689

chris-tsiaousis-hpc

Nice PR. I wonder what performance implication will the second K iter for B have (from what I've understand so far this was for cache-hits on A).
Great work so far @EnricoDeg ! The PR as a whole became rather complex but at least github has this functionality that you can view the merged PR as one commit.
I also have some other questions and some suggestions. See my comments:

…3689

chris-tsiaousis-hpc

It's up to a CI pass for me now. Well done!

[CK_TILE] Extend support of mix precision microscaling BQuant (#4267) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Proposed changes Supported types combinations using BQuant=e8m0: - A=bf16 - B=bf16,bf8,fp4 Summary: - remove usage of `pk_fp4_raw_t`: consistent with other implementations and avoid taking into account of the packed size explicitly. In general, the raw type should not be used because CK Tile internally takes care of the PackedSize, so using the raw type adds unnecessary complexity to the implementation - handle microscaling by checking for `e8m0` type for BQuant (previous implementation was inconsistent) - add support for scaling instructions in `DequantPack8` - mx pipeline: - extend existing pipeline to support different B types - add support to scale and cast before writing to LDS or after reading from LDS (this can be defined in the `Problem` by the user) - block gemm: - mx pipeline is now using block gemm BQuant - block gemm BQuant can now load from LDS and apply scale and then call block gemm universal operator. This adds new functionalities and remove code duplication - warp gemm: - add case to support 128bit ds_read/write for both A and B when A=16bit and B=8bit - add examples and tests: note that some tests for bf16/fp4 already existed but were removed during previous tests refactoring. I added them again and other relevant tests for new types combinations ## 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

) ## Proposed changes Supported types combinations using BQuant=e8m0: - A=bf16 - B=bf16,bf8,fp4 Summary: - remove usage of `pk_fp4_raw_t`: consistent with other implementations and avoid taking into account of the packed size explicitly. In general, the raw type should not be used because CK Tile internally takes care of the PackedSize, so using the raw type adds unnecessary complexity to the implementation - handle microscaling by checking for `e8m0` type for BQuant (previous implementation was inconsistent) - add support for scaling instructions in `DequantPack8` - mx pipeline: - extend existing pipeline to support different B types - add support to scale and cast before writing to LDS or after reading from LDS (this can be defined in the `Problem` by the user) - block gemm: - mx pipeline is now using block gemm BQuant - block gemm BQuant can now load from LDS and apply scale and then call block gemm universal operator. This adds new functionalities and remove code duplication - warp gemm: - add case to support 128bit ds_read/write for both A and B when A=16bit and B=8bit - add examples and tests: note that some tests for bf16/fp4 already existed but were removed during previous tests refactoring. I added them again and other relevant tests for new types combinations ## 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 --- 🔁 Imported from [ROCm/composable_kernel#3689](ROCm/composable_kernel#3689) 🧑‍💻 Originally authored by @EnricoDeg --------- Co-authored-by: Enrico Degregori <enrico@streamhpc.com> Co-authored-by: systems-assistant[bot] <systems-assistant[bot]@users.noreply.github.com> Co-authored-by: Thomas Ning <Thomas.Ning@amd.com> Co-authored-by: Enrico Degregori <73224202+EnricoDeg@users.noreply.github.com> Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com>

@EnricoDeg

[CK_TILE] Extend support of mix precision microscaling BQuant (#4267) ## Proposed changes Supported types combinations using BQuant=e8m0: - A=bf16 - B=bf16,bf8,fp4 Summary: - remove usage of `pk_fp4_raw_t`: consistent with other implementations and avoid taking into account of the packed size explicitly. In general, the raw type should not be used because CK Tile internally takes care of the PackedSize, so using the raw type adds unnecessary complexity to the implementation - handle microscaling by checking for `e8m0` type for BQuant (previous implementation was inconsistent) - add support for scaling instructions in `DequantPack8` - mx pipeline: - extend existing pipeline to support different B types - add support to scale and cast before writing to LDS or after reading from LDS (this can be defined in the `Problem` by the user) - block gemm: - mx pipeline is now using block gemm BQuant - block gemm BQuant can now load from LDS and apply scale and then call block gemm universal operator. This adds new functionalities and remove code duplication - warp gemm: - add case to support 128bit ds_read/write for both A and B when A=16bit and B=8bit - add examples and tests: note that some tests for bf16/fp4 already existed but were removed during previous tests refactoring. I added them again and other relevant tests for new types combinations ## 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 --- 🔁 Imported from [#3689](#3689) 🧑‍💻 Originally authored by @EnricoDeg --------- Co-authored-by: Enrico Degregori <enrico@streamhpc.com> Co-authored-by: systems-assistant[bot] <systems-assistant[bot]@users.noreply.github.com> Co-authored-by: Thomas Ning <Thomas.Ning@amd.com> Co-authored-by: Enrico Degregori <73224202+EnricoDeg@users.noreply.github.com> Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com>

EnricoDeg and others added 17 commits February 3, 2026 17:20

Generalize implementation to support different types

8684a67

Add tests

e79e609

Use pk cvt instruction bf8 to bf16 on gfx950

aecc8b0

Fix fp4 * scale -> bf16 to use pk instruction on gfx950

67c72d3

Fix vectorsize buffer load for 16/8 bit case

45063ac

Fix LDS read/write for 16/8 bit case

02b81bd

Both A and B are using 128 bit instructions

Refactor pipeline

e2c1dcd

Support both scale before writing to LDS or scale after reading from LDS

Naming convention examples mx blockscale

16491f0

Rename pipeline (not limited to fp4 any more)

e907e1b

Clean up tile_elementwise (casting not needed with new approach)

33f4f87

Clean up pipeline

afd3d2b

Finalize cleanup

b6a66c1

Fix rebase issues

3dde9eb

Fix usage of Problem static member BCastPolicy

0655316

Move static constexpr to structs

17e8f2f

Fix initialization examples

c89dfbf

Merge commit 'c89dfbf1534a730e537de711a9bf084a9e2e8fa1' into import/d…

40c66ea

…evelop/ROCm_composable_kernel/pr-3689

assistant-librarian Bot requested review from ThomasNing, afagaj, andriy-ca, aosewski, asleepzzz, bartekxk, carlushuang, geyyer, illsilin, poyenc, qianfengz and tenpercent as code owners February 3, 2026 21:46

assistant-librarian Bot added the imported pr label Feb 3, 2026

illsilin approved these changes Feb 10, 2026

View reviewed changes

EnricoDeg self-assigned this Feb 10, 2026

Merge branch 'develop' into import/develop/ROCm_composable_kernel/pr-…

2242bb0

…3689

ThomasNing reviewed Feb 11, 2026

View reviewed changes

EnricoDeg and others added 4 commits February 12, 2026 08:07

More explicit builtin arguments

af09c9b

Rename pipeline files and structs

05d22ce

Merge branch 'develop' into import/develop/ROCm_composable_kernel/pr-…

5d78f76

…3689

Merge branch 'develop' into import/develop/ROCm_composable_kernel/pr-…

4fd0390

…3689

ThomasNing enabled auto-merge (squash) February 12, 2026 18:33

EnricoDeg and others added 3 commits February 13, 2026 08:27

Merge branch 'develop' into import/develop/ROCm_composable_kernel/pr-…

cc57bcf

…3689

Merge branch 'develop' into import/develop/ROCm_composable_kernel/pr-…

9f34a7b

…3689

EnricoDeg marked this pull request as draft February 18, 2026 13:41

auto-merge was automatically disabled February 18, 2026 13:41
Pull request was converted to draft

Merge branch 'develop' into import/develop/ROCm_composable_kernel/pr-…

e23b35c

…3689

chris-tsiaousis-hpc requested changes Feb 19, 2026

View reviewed changes

EnricoDeg added 5 commits February 23, 2026 08:07

Merge branch 'develop' into import/develop/ROCm_composable_kernel/pr-…

bf67114

…3689

Fix condition OverrideBDataType

0f7502f

Refactor conversion dequant

a7d5ad6

Fix comment

93df781

Fix handling of AttrNumAccess

ac33944

chris-tsiaousis-hpc approved these changes Feb 23, 2026

View reviewed changes

EnricoDeg marked this pull request as ready for review February 23, 2026 13:40

ThomasNing merged commit 3c5d95e into develop Feb 24, 2026
50 checks passed

ThomasNing deleted the import/develop/ROCm_composable_kernel/pr-3689 branch February 24, 2026 17:55

EnricoDeg mentioned this pull request Apr 1, 2026

[CK Tile] Add transposed tile load implementation, and tests for load_and_convert_tile #5510

Merged

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[CK_TILE] Extend support of mix precision microscaling BQuant - #4267

[CK_TILE] Extend support of mix precision microscaling BQuant#4267
ThomasNing merged 37 commits into
developfrom
import/develop/ROCm_composable_kernel/pr-3689

assistant-librarian Bot commented Feb 3, 2026

Uh oh!

ThomasNing left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chris-tsiaousis-hpc left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chris-tsiaousis-hpc left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

assistant-librarian Bot commented Feb 3, 2026

Proposed changes

Checklist

Discussion

Uh oh!

ThomasNing left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chris-tsiaousis-hpc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chris-tsiaousis-hpc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants