[MHA] Integrate ck mha fp8 solver into miopen #3014

bghimireamd · 2024-06-03T13:11:16Z

CK's fp8 col major solver integration into MIOpen

CAHEK7

some minor suggestions

src/solver/mha/mha_ck_solver_forward.cpp

…ntegration

CAHEK7 · 2024-08-05T17:16:37Z

requirements.txt

@@ -7,5 +7,5 @@ nlohmann/json@v3.11.2 -DJSON_MultipleHeaders=ON -DJSON_BuildTests=Off
 ROCm/FunctionalPlus@v0.2.18-p0
 ROCm/eigen@3.4.0
 ROCm/frugally-deep@9683d557eb672ee2304f80f6682c51242d748a50
-ROCm/composable_kernel@15baccf2ecad4fb3498b8acb6bbf58fb5359c7a5 -DCMAKE_BUILD_TYPE=Release -DINSTANCES_ONLY=ON
+ROCm/composable_kernel@1dd9875c9ac783cabc8a536e0802b58d43b6b107 -DCMAKE_BUILD_TYPE=Release -DINSTANCES_ONLY=ON


Can this be aligned with #3142 and #3181?

Yes, I also recommend picking CK commit hashes from its amd-develop branch.

CAHEK7 · 2024-08-05T17:18:07Z

src/include/miopen/mha/problem_description.hpp

+        const auto& fptr = GetDescsForward();
+        if(fptr.oDesc.GetType() == miopenFloat8)
+        {
+            return true;
+        }
+        else
+        {
+            return false;
+        }


Suggested change

const auto& fptr = GetDescsForward();

if(fptr.oDesc.GetType() == miopenFloat8)

{

return true;

}

else

{

return false;

}

return GetDescsForward().oDesc.GetType() == miopenFloat8;

But I'm not sure that it has to be here at all.

I was following convolution's problem description pattern.

Then my suggestion matches this pattern perfectly.

CAHEK7 · 2024-08-05T17:23:32Z

src/solver/mha/mha_ck_solver_forward.cpp

+    const std::string name = context.GetStream().GetDeviceName();
+    if(!(name == "gfx940" || name == "gfx941" || name == "gfx942"))
+        return false;


StartsWith(handle.GetDeviceName(), "gfx103")

Suggested change

const std::string name = context.GetStream().GetDeviceName();

if(!(name == "gfx940" || name == "gfx941" || name == "gfx942"))

return false;

if(!StartsWith(handle.GetDeviceName(), "gfx94"))

return false;

CAHEK7 · 2024-08-05T17:25:48Z

src/solver/mha/mha_ck_solver_forward.cpp

+    if(!problem.IsFFp8()) // forward mha fp8
+        return false;


Probably all the tensors should be explicitly checked.

CAHEK7 · 2024-08-05T17:27:07Z

src/solver/mha/mha_ck_solver_forward.cpp

+    ::miopen::mha::MhaInputDescsForward mha_des = problem.GetDescsForward();
+    const auto& lens                            = mha_des.qDesc.GetLengths();
+    auto [N, H, S, D]                           = std::tie(lens[0], lens[1], lens[2], lens[3]);


Suggested change

::miopen::mha::MhaInputDescsForward mha_des = problem.GetDescsForward();

const auto& lens = mha_des.qDesc.GetLengths();

auto [N, H, S, D] = std::tie(lens[0], lens[1], lens[2], lens[3]);

auto [N, H, S, D] = miopen::tien<4>(problem.GetDescsForward().qDesc.GetLengths());

CAHEK7 · 2024-08-05T17:28:05Z

src/solver/mha/mha_ck_solver_forward.cpp

+    if(D <= 256 && S % 128 == 0 && D % 64 == 0)
+        return true;
+    return false;


Suggested change

if(D <= 256 && S % 128 == 0 && D % 64 == 0)

return true;

return false;

return (D <= 256 && S % 128 == 0 && D % 64 == 0);

CAHEK7 · 2024-08-05T18:00:17Z

test/gtest/gpu_mha_ck_forward.cpp

+        auto [n, h, s, d, drop] = GetParam();
+        Handle& handle          = get_handle();
+
+        if((drop > 0.0f) && (s % handle.GetWavefrontWidth() != 0))


As far as I know, CK does not support dropout at all.

CAHEK7 · 2024-08-05T18:00:58Z

test/gtest/gpu_mha_ck_forward.cpp

+        InitTensor(miopenTensorMhaV, std::move(v.mTensor));
+
+        float s_scale = test::cpu::GetF8Scaling(1.0);
+        // clang-tidy complains about the same expression on both sides of "/": 1.f / 1.f


Suggested change

// clang-tidy complains about the same expression on both sides of "/": 1.f / 1.f

CAHEK7 · 2024-08-05T18:07:59Z

test/gtest/gpu_mha_ck_forward.cpp

+    float amaxO_ref;
+
+    bool is_ck_solver             = false;
+    float ck_fp8_solver_threshold = 0.015;


This is too high threshold, MIOpen fp8 algorithm fits into 0.0002 which is 75 times smaller.
Could you double-check that the values are similar, because rms error is not sensible to permutations and some significant outliers.

CAHEK7 · 2024-08-05T18:08:31Z

test/gtest/gpu_mha_ck_forward.cpp

+class Test_Fwd_Mha_F32 : public Test_Fwd_Mha<float>
+{
+};


As far as I know, CK does not support FP32.

CAHEK7 · 2024-08-05T18:10:29Z

test/gtest/mha_helper.hpp

+    else
+    {
+        // ck solver
+        ScaleMult(atten_heads_fp32, GetF8Scaling(aMax_O), multi_head_attention_fp8);


aMax_O must not be used here, it's an output parameter and will be passed as scale_o into the next iteration.

junliume · 2024-08-21T06:19:07Z

@bghimireamd do you wish to follow up with this PR?

bghimireamd · 2024-09-12T14:29:38Z

Current fp8 V2 implementation is very limited. We will place close this PR since CK now implementing fp8 v3. We will have new PR for the new CK fp8 v3 integration.

initial comment for attemting to integrate ck mha fp8 solver

2b9bb7b

bghimireamd changed the title ~~initial comment for attemting to integrate ck mha fp8 solver~~ Integrate ck mha fp8 solver into miopen Jun 3, 2024

bghimireamd changed the title ~~Integrate ck mha fp8 solver into miopen~~ [DRAFT][MHA] Integrate ck mha fp8 solver into miopen Jun 3, 2024

bghimireamd added 2 commits June 9, 2024 14:43

compiling now

3140f10

return the solution

0aeb5b3

CAHEK7 reviewed Jun 10, 2024

View reviewed changes

src/solver/mha/mha_ck_solver_forward.cpp Outdated Show resolved Hide resolved

src/solver/mha/mha_ck_solver_forward.cpp Outdated Show resolved Hide resolved

bghimireamd added 6 commits July 1, 2024 18:07

test passing now

4b2687b

in middle

ad00148

both cpu and gpu results from ck solvers are matching

c7c72b6

fix review comments

804d2a9

add workspace in ck solver

eb9699d

fix merge conflict

a7f1824

bghimireamd marked this pull request as ready for review July 19, 2024 18:49

bghimireamd requested review from JehandadKhan and junliume as code owners July 19, 2024 18:49

bghimireamd added 2 commits July 19, 2024 18:56

minor cleanups

5cd58d4

uncomment test

c5e35be

bghimireamd changed the title ~~[DRAFT][MHA] Integrate ck mha fp8 solver into miopen~~ [MHA] Integrate ck mha fp8 solver into miopen Jul 19, 2024

bghimireamd marked this pull request as draft July 21, 2024 15:13

create sep test for ck mha and miopen mha

2568561

bghimireamd marked this pull request as ready for review July 23, 2024 19:22

bghimireamd added 2 commits July 23, 2024 19:22

Merge branch 'develop' of github.com:ROCm/MIOpen into bg/ck_mha_fp8_i…

759c640

…ntegration

update CK commit hash

f3922e0

bghimireamd requested a review from CAHEK7 July 24, 2024 14:05

bghimireamd added 6 commits July 24, 2024 14:10

add missing test file

c8172a3

fix commit hash

a53482d

update commit hash

15c0b52

update ck commit hash

5accfe7

update new ck commit hash

52f1312

update ck hash to have ck_tile

f1d1862

bghimireamd added 8 commits August 1, 2024 01:46

update ck commit hash

f7c25c7

ck commit hash update

6cc7b92

ck_tile correct path

9b01b10

bg/ck_mha_fp8_integration : fix compile error

cd163e6

gtest pass

776c103

fix tidy issue

3e23eb0

fix tidy issue

9b4b11f

profiling mha ck solver

9dc7534

bghimireamd added the TESTING_CI_PASSED label Aug 4, 2024

minor fixes

a90e188

bghimireamd requested a review from carlushuang August 5, 2024 12:42

CAHEK7 reviewed Aug 5, 2024

View reviewed changes

CAHEK7 requested changes Aug 5, 2024

View reviewed changes

bghimireamd closed this Sep 12, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MHA] Integrate ck mha fp8 solver into miopen #3014

[MHA] Integrate ck mha fp8 solver into miopen #3014

bghimireamd commented Jun 3, 2024 •

edited

Loading

CAHEK7 left a comment

CAHEK7 Aug 5, 2024 •

edited

Loading

junliume Aug 7, 2024

CAHEK7 Aug 5, 2024

bghimireamd Aug 21, 2024

CAHEK7 Aug 21, 2024

CAHEK7 Aug 5, 2024

CAHEK7 Aug 5, 2024

CAHEK7 Aug 5, 2024

CAHEK7 Aug 5, 2024

CAHEK7 Aug 5, 2024

CAHEK7 Aug 5, 2024

CAHEK7 Aug 5, 2024

CAHEK7 Aug 5, 2024

CAHEK7 Aug 5, 2024

junliume commented Aug 21, 2024

bghimireamd commented Sep 12, 2024 •

edited

Loading

[MHA] Integrate ck mha fp8 solver into miopen #3014

[MHA] Integrate ck mha fp8 solver into miopen #3014

Conversation

bghimireamd commented Jun 3, 2024 • edited Loading

CAHEK7 left a comment

Choose a reason for hiding this comment

CAHEK7 Aug 5, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

junliume commented Aug 21, 2024

bghimireamd commented Sep 12, 2024 • edited Loading

bghimireamd commented Jun 3, 2024 •

edited

Loading

CAHEK7 Aug 5, 2024 •

edited

Loading

bghimireamd commented Sep 12, 2024 •

edited

Loading