Add zero dim tensor check when using flash_attention

[ranzhejiang]

The cuda or triton kernel can not support this case: dimensions of size is 0, but traditional SDPA is implemented using PyTorch's tensor operations, which have robust support for tensors with dimensions of size 0. We need to add this check and error tips for developers when using flash_attention. Related issue is in deepspeedai/DeepSpeed#7275

[Rocketknight1]

cc @ArthurZucker for FA2

[ArthurZucker]

Thanks, I'd rather we write 3 explicit checks as only query, key and value need this check. Maybe even just testing key?

[ranzhejiang]

Thanks, I'd rather we write 3 explicit checks as only query, key and value need this check. Maybe even just testing key?

Do you mean we don't need write a function to do this check, just call explicit check directly? if so, I agree with you, I can write 3 explicit checks for query, key and value, or only for key, and give users info which shape is wrong.

[ArthurZucker]

Yeah, I mean only checking q should be enough as well no?

[ranzhejiang]

Yeah, I mean only checking q should be enough as well no?

Yes, I agree with you and have change my code, thanks for review

[ranzhejiang]

@ArthurZucker Hi ArthurZucker, I have changed my code following your advice, can you help review this PR ? Thanks

[ArthurZucker]

sorry I was out for holidays