- Fix Half/BFloat16 dtype mismatch in txt_in layer - Add distill_vec parameter to forward function signature