Support hdimQK != hdimV backward

[shcho1118]

Recently, there are cases where hdimQK >= hdimV hdimQK != hdimV, such as DeepSeek (#1487).
So we worked on a simple trick (rounding hdimQK and hdimV to the same value) to support that configuration in the backward kernel as well.
Of course, it would be best to tile K and V with different rounding values (which we're working on internally), but this implementation allows us to skip explicit padding & unpadding hdimV, resulting in a performance improvement of about 15%.

This feature was developed for internal use at @kakao.

[tridao]

Thank you! Could you rebase on the latest in main and then I'll merge?

[shcho1118]

@tridao
We've done the rebase as you requested, and I think we're ready to merge.

[tridao]

The hdim dispatching should be based on max(params.d, params.d_v) instead of based on params.d?
E.g. if d = 128, d_v = 192, then it should dispatch to kHeadDim=192 instead of kHeadDim=128?

[shcho1118]

I was initially thinking about the case where d = 192, d_v = 128 like deepseek, but max(params.d, params.d_v) would be more universal.

[shcho1118]

@tridao
I took your advice and extended the support a bit more, now if max(params.d, params.dv) <= 256 it should work fine.

[ehuaa]

Hi @shcho1118 @tridao , I noticed that this pr has implemented hdimQK != hdimV for backward for both sm80 and sm90 arch, but for forward (mha_forward function)

flash-attention/hopper/flash_api.cpp

Line 759 in a9a3170

TORCH_CHECK(dprops->major == 9, "Only Hopper supports different V headdim");

TORCH_CHECK(dprops->major == 9, "Only Hopper supports different V headdim");
there's a if-statement that only hopper arch can run deepseek series models in foward function
Do you have any ideas of how to bypass this if-statement and run deepseek with sm80 arch for forward function

[shcho1118]

@ehuaa
After I finished implementing sm80 backward, I discovered that if-statement, so sm80 backward kernel is not fully validated in this PR.
If you want to enable the same functionality in sm80 forward kernel, the fastest way to do it is to touch rounding hdim like this PR.

[ehuaa]

@ehuaa After I finished implementing sm80 backward, I discovered that if-statement, so sm80 backward kernel is not fully validated in this PR. If you want to enable the same functionality in sm80 forward kernel, the fastest way to do it is to touch rounding hdim like this PR.

Thanks for your quick reply, i'll try it later.

[Graham1025]

Recently, there are cases where hdimQK >= hdimV hdimQK != hdimV, such as DeepSeek (#1487). So we worked on a simple trick (rounding hdimQK and hdimV to the same value) to support that configuration in the backward kernel as well. Of course, it would be best to tile K and V with different rounding values (which we're working on internally), but this implementation allows us to skip explicit padding & unpadding hdimV, resulting in a performance improvement of about 15%.

This feature was developed for internal use at @kakao.
Thanks for your great work! So does it use TMA oob feature for (rounding hdimQK and hdimV to the same value)?

[shcho1118]

@Graham1025
Yes, for sm90, it was already using that feature to do rounding, so I just modified it slightly.