[CK_TILE] Support more layouts for BQuant GEMM #3349
Conversation
…w_layouts_datatypes
…w_layouts_datatypes
…w_layouts_datatypes
samremes
requested review from
ThomasNing,
afagaj,
andriy-ca,
aosewski,
asleepzzz,
bartekxk,
carlushuang,
cgmillette,
coderfeli,
geyyer,
illsilin,
poyenc,
qianfengz and
tenpercent
as code owners
ThomasNing
left a comment
ThomasNing
left a comment
There was a problem hiding this comment.
Thank you for the contributions. LGTM except the above comments.
| static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>); | ||
| // PreshuffleQuant currently assumes ColumnMajor layout | ||
| // For RowMajor, the preshuffle logic would need adjustment | ||
| static_assert(std::is_same_v<BQLayout, tensor_layout::gemm::ColumnMajor>, |
There was a problem hiding this comment.
@amd-khushbu Please not this part as our discussion yesterday. With the preshuffle quant our tensor layout should actually only be RowMajor. We could modify that in your PR.
There was a problem hiding this comment.
I'll revert the comment here to avoid creating any confusion.
| decltype(make_static_distributed_tensor<ADataType>(ABlockTileDistr{})); | ||
| using BBlockTile = | ||
| decltype(make_static_distributed_tensor<BDataType>(BBlockTileDistr{})); | ||
| decltype(make_static_distributed_tensor<ADataType>(BBlockTileDistr{})); |
There was a problem hiding this comment.
Here should be BDataType?
There was a problem hiding this comment.
I think this should be the non-pkint4 datatype, which should be the same as ADataType, but maybe it should be made more clear somehow. I'll create an alias for the type name, so it avoids confusion.
| b_dram_block_window_tmp, | ||
| [](const BDataType& b) { return b; }, | ||
| // Note: BDataType gets converted to ADataType during loading | ||
| [](const ADataType& b) { return b; }, |
There was a problem hiding this comment.
That condition should only happen when the BDataType is pk_int4 right?
Proposed changes
Enables more layouts for BQuant GEMM. Row/Col for A is tested, and support for Row/Col for B and BQ is added.
Checklist
Please put an
xinto 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.clang-formaton all changed filesDiscussion
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