Bf16*fp4 gemm #2801
Bf16*fp4 gemm #2801
Conversation
eliotwang
requested review from
a team,
ThomasNing,
afagaj,
andriy-ca,
aosewski,
asleepzzz,
bartekxk,
carlushuang,
coderfeli,
ddembeckAMD,
geyyer,
illsilin,
poyenc,
qianfengz,
shumway,
tenpercent and
vidyasagar-amd
as code owners
ThomasNing
left a comment
ThomasNing
left a comment
There was a problem hiding this comment.
Could we add the gtest for the new developed kernel?
| #include "ck_tile/host.hpp" | ||
| #include "gemm_utils.hpp" | ||
|
|
||
| template <typename GemmConfig, |
There was a problem hiding this comment.
We should also put the mx gemm into the example of blockscale gemm and share the util and example.inc code.
There was a problem hiding this comment.
So, for reusing the code, my understanding is that we need to add our own mx_gemm.cpp entry interface within example/38_* instead of defining a new example/45_* like we are doing now?
There was a problem hiding this comment.
Yes, we do not need to create a mxfp4_gemm example operator. We should just have a .cpp file under example/38_*/mx_gemm.cpp. The datatype should also be a configuration to that example..
| make_tuple(kargs.stride_B, 1), | ||
| number<GemmPipeline::GetVectorSizeB()>{}, | ||
| number<1>{}); | ||
| if constexpr(std::is_same_v<BDataType, pk_fp4_raw_t>) |
There was a problem hiding this comment.
Why in fp4 we will transposed in the data type comparing to other data types?
There was a problem hiding this comment.
In the original gemm_quant kernel, we noticed that for the tensor_view definitions of B and Bq, B's shape is taken as (N, K), while Bq is (K, N). We don't quite understand why Bq needs to be transposed here, as in our implementation both B and Bq are defined as (N, K).
There was a problem hiding this comment.
I see. After reviewing the solution, we could reunify into one version without the transpose. I will create a PR to unify that soon, so we do not need that unnecessary transpose branch.
|
Hi @eliotwang, please resolve conflicts and sync branch to latest develop in order to proceed! Thanks! |
|
@eliotwang LGTM overall. Please add the unit test. |
We have added unit tests for bf16_mxfp4_gemm in the test/ck_tile/gemm_block_scale/ directory. Please help review it. |
| { | ||
| using ComputeType = | ||
| std::conditional_t<sizeof(ADataType_) < sizeof(BDataType_), ADataType_, BDataType_>; | ||
| // using ComputeType = |
There was a problem hiding this comment.
Has been updated. Please help review it.
|
@eliotwang LGTM, we could do the last iteration of the merging after the PR #3245 merged to the develop. Thanks! cc. @CongMa13 |
7 participants
Proposed changes
Added an example of bf16*fp4 gemm, where fp4 and fp4_scale are in uint8 data format. In the pipeline, matrix B(fp4) will be dequantized to bf16 before performing multiplication operations.
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