fix: align Gemma dtype before RoPE

[wilburx813]

Summary

This issue was observed when running Pi0.5 inference through RPent. Originally reported in RLinf/RPent#131. The model loaded successfully, but inference failed with:

RuntimeError: expected scalar type BFloat16 but found Float

This regression was introduced while porting the custom Gemma implementation from Transformers 4.53 to 4.57. The existing BF16 cast was moved after RoPE construction, causing RoPE to remain FP32 while the attention values and prefix KV cache were BF16.

The resulting dtype mismatch caused the attention matmul to fail.

Fix

Restore the original ordering by moving the existing BF16 cast before RoPE construction.

This keeps the RoPE and attention tensors in BF16 without adding new casting logic.

Verification

• Before: the KV-cache inference path fails with an FP32/BF16 mismatch.
• After: the same path returns finite BF16 output.
• Real Pi0.5 checkpoint inference on an RTX 4090 returns finite actions with shape (1, 5, 7).