diff --git a/tensorflow/core/kernels/gpu_prim.h b/tensorflow/core/kernels/gpu_prim.h
index bcc873bfb03180..712663cfd88bbe 100644
--- a/tensorflow/core/kernels/gpu_prim.h
+++ b/tensorflow/core/kernels/gpu_prim.h
@@ -82,8 +82,8 @@ namespace gpuprim = ::hipcub;
 
 // Required for sorting Eigen::half and bfloat16.
 namespace rocprim {
+#if (TF_ROCM_VERSION >= 50200 && TF_ROCM_VERSION < 70000)
 namespace detail {
-#if (TF_ROCM_VERSION >= 50200)
 template <>
 struct float_bit_mask<Eigen::half> {
   static constexpr uint16_t sign_bit = 0x8000;
@@ -99,7 +99,27 @@ struct float_bit_mask<Eigen::bfloat16> {
   static constexpr uint16_t mantissa = 0x007F;
   using bit_type = uint16_t;
 };
+}; // namespace detail
+
+#else
+namespace traits {
+template<>
+struct rocprim::traits::define<Eigen::half> {
+  using float_bit_mask = rocprim::traits::float_bit_mask::values<uint16_t, 0x8000, 0x7C00, 0x03FF>;
+  using is_arithmetic = rocprim::traits::is_arithmetic::values<true>;
+  using number_format = rocprim::traits::number_format::values<traits::number_format::kind::floating_point_type>;
+};
+
+template<>
+struct rocprim::traits::define<tsl::bfloat16> {
+  using float_bit_mask = rocprim::traits::float_bit_mask::values<uint16_t, 0x8000, 0x7F80, 0x007F>;
+  using is_arithmetic = rocprim::traits::is_arithmetic::values<true>;
+  using number_format = rocprim::traits::number_format::values<traits::number_format::kind::floating_point_type>;
+};
+}; // namespace traits
 #endif
+#if (TF_ROCM_VERSION < 70000)
+namespace detail {
 template <>
 struct radix_key_codec_base<Eigen::half>
     : radix_key_codec_floating<Eigen::half, uint16_t> {};
@@ -107,6 +127,7 @@ template <>
 struct radix_key_codec_base<tensorflow::bfloat16>
     : radix_key_codec_floating<tensorflow::bfloat16, uint16_t> {};
 };  // namespace detail
+#endif
 };  // namespace rocprim
 
 #endif  // TENSORFLOW_USE_ROCM