fix: preserve input dtype instead of hardcoding float16 in attention layers

[Mr-Neutr0n]

Summary

• Replace hardcoded .half() calls with .to(x.dtype) in hydit/modules/attn_layers.py to preserve the original input tensor's dtype
• Fixes silent precision loss when training with bfloat16 (common on A100/H100 GPUs), where .half() forces a downcast from bfloat16 to float16
• Affected locations: FlashSelfMHAModified.forward() (2 calls), FlashCrossMHAModified.forward() (3 calls), and CrossAttention.forward() IP-adapter path (1 call)

Problem

Several places in the attention layer implementations hardcode .half() after QK normalization and softmax operations. This forces the tensor dtype to float16 regardless of the actual training precision. When running in bfloat16 mode, this silently downcasts tensors from bfloat16 to float16, causing:

• Precision loss due to bfloat16 having a larger dynamic range than float16
• Potential numerical instability during training
• Inconsistent dtype handling within the forward pass

Fix

Replace all .half() calls with .to(x.dtype), which dynamically preserves whatever dtype the input tensor x uses. This is a backward-compatible change -- when running in float16 mode, x.dtype will be torch.float16 and behavior is identical to before.

Test plan

• Verify existing float16 training produces identical results (.to(x.dtype) == .half() when input is float16)
• Verify bfloat16 training no longer silently downcasts to float16
• Verify inference outputs are unchanged for float16 models

[tencent-adm]

Thank you for your submission, we really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[Mr-Neutr0n]

I have read the CLA Document and I hereby sign the CLA

[Mr-Neutr0n]

recheck