Update schema in test_flash3_bw_compatibility

[guilhermeleobas]

xref: #1769 (comment)

Pull request #1769 breaks backward compatibility of the backward function as the C++ mha_bwd function returns some of its inputs (dq, dk, dv).

cc: @janeyx99

[janeyx99]

Thanks, I do think the longterm solution if we were to fix the test is to update the schema.

[guilhermeleobas]

Isn't the test testing for backwards compatibility with Flash Attention 2?

[janeyx99]

@guilhermeleobas do you mean this schema is from FA2? https://github.com/Dao-AILab/flash-attention/blob/main/hopper/test_flash_attn.py#L1244-L1250

I wasn't aware they were tied.

[guilhermeleobas]

Let me double-check that info, if FA3 schema is backward compatible with FA2.

I'm not sure if that's feasible, but I can try to return empty tensors for dq, dk, dv in the backward function instead of breaking backward compatibility.

flash-attention/hopper/flash_api_stable.cpp

Line 1645 in 13696f2

return { softmax_d, softmax_lse_log2, dq_accum, dk_accum, dv_accum };

[guilhermeleobas]

I wasn't aware they were tied.

I think you're right, FA2 and FA3 schemas are not tied.

[janeyx99]

If we're okay BC breaking originally, this change now LGTM as long as the test passes locally.

I'm not sure if the alternative where we undo the BC break is better because it's already been a month or so and people have been building wheels with the current FA3. So changing it back will BC break those people 😬