diff --git a/fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu b/fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu
index 781c159d8a..b976de4f33 100644
--- a/fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu
+++ b/fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu
@@ -1257,17 +1257,15 @@ Tensor {{ embedding_cuda_op }}(
                     const static auto use_hip_kernel = fbgemm_gpu::config::is_feature_enabled(fbgemm_gpu::config::FeatureGateName::TBE_ROCM_HIP_BACKWARD_KERNEL);
 
                     constexpr bool supported_weights_type = std::is_same_v<emb_t, float> || std::is_same_v<emb_t, at::Half>;
-                    constexpr bool supported_grad_type = std::is_same_v<grad_t, float> || std::is_same_v<grad_t, at::Half>;
+                    constexpr bool supported_grad_type = std::is_same_v<grad_t, float> || std::is_same_v<grad_t, at::Half> || std::is_same_v<grad_t, at::BFloat16>;
                     const bool cached = uvm_weights.numel() > 0 || lxu_cache_weights.numel() > 0;
 
-                    constexpr bool same_precision = std::is_same_v<emb_t, grad_t>;
-
                     // The optimized HIP backward kernel uses warpSize-64-only
                     // intrinsics; it must only be dispatched on warpSize 64
                     // devices. warpSize 32 devices fall through to the generic
                     // shuffle-based backward kernel.
 
-                    if (use_hip_kernel && at::cuda::warp_size() == 64 && !mixed_D && !cached && supported_weights_type && supported_grad_type && same_precision && rocm::is_supported_cdna())
+                    if (use_hip_kernel && at::cuda::warp_size() == 64 && !mixed_D && !cached && supported_weights_type && supported_grad_type && rocm::is_supported_cdna())
                     {
                         constexpr int segments_per_workgroup = 4;
                         {%- for kDimSize in [64, 128, 160, 192, 256, 320] %}
diff --git a/fbgemm_gpu/include/fbgemm_gpu/rocm/split_embeddings_common.h b/fbgemm_gpu/include/fbgemm_gpu/rocm/split_embeddings_common.h
index d7d6954d00..2065e590e7 100644
--- a/fbgemm_gpu/include/fbgemm_gpu/rocm/split_embeddings_common.h
+++ b/fbgemm_gpu/include/fbgemm_gpu/rocm/split_embeddings_common.h
@@ -88,6 +88,12 @@ __device__ float llvm_amdgcn_raw_buffer_load_fp32(
     int32_t soffset = 0,
     int32_t glc_slc = 0) __asm("llvm.amdgcn.raw.buffer.load.f32");
 
+__device__ floatx2_t llvm_amdgcn_raw_buffer_load_fp32x2(
+    int32x4_t srsrc,
+    int32_t voffset,
+    int32_t soffset = 0,
+    int32_t glc_slc = 0) __asm("llvm.amdgcn.raw.buffer.load.v2f32");
+
 __device__ half2 llvm_amdgcn_raw_buffer_load_fp16x2(
     int32x4_t srsrc,
     int32_t voffset,
@@ -160,81 +166,27 @@ struct load_row_per_warp {
   }
 };
 
-template <typename index_t>
-struct load_row_per_warp<half, 64, index_t> {
-  static __device__ void
-  run(half* emb_data, index_t row_index, const half* p_emb_table, int lane_id) {
-    int32x4_t emb_res =
-        amdgcn_make_buffer_resource(p_emb_table + row_index * 64);
-    emb_data[0] =
-        llvm_amdgcn_raw_buffer_load_fp16(emb_res, lane_id * sizeof(half));
-  }
-};
-
-template <typename index_t>
-struct load_row_per_warp<half, 128, index_t> {
-  static __device__ void
-  run(half* emb_data, index_t row_index, const half* p_emb_table, int lane_id) {
-    int32x4_t emb_res =
-        amdgcn_make_buffer_resource(p_emb_table + row_index * 128);
-    *reinterpret_cast<half2*>(emb_data) =
-        llvm_amdgcn_raw_buffer_load_fp16x2(emb_res, lane_id * sizeof(half2));
-  }
-};
-
-template <typename index_t>
-struct load_row_per_warp<half, 160, index_t> {
-  static __device__ void
-  run(half* emb_data, index_t row_index, const half* p_emb_table, int lane_id) {
-    int32x4_t emb_res = amdgcn_make_buffer_resource(
-        p_emb_table + row_index * 160, sizeof(half) * 160);
-    *reinterpret_cast<half2*>(emb_data) =
-        llvm_amdgcn_raw_buffer_load_fp16x2(emb_res, lane_id * sizeof(half2));
-    emb_data[2] = llvm_amdgcn_raw_buffer_load_fp16(
-        emb_res, (lane_id + 128) * sizeof(half));
-  }
-};
-
-template <typename index_t>
-struct load_row_per_warp<half, 192, index_t> {
-  static __device__ void
-  run(half* emb_data, index_t row_index, const half* p_emb_table, int lane_id) {
-    int32x4_t emb_res =
-        amdgcn_make_buffer_resource(p_emb_table + row_index * 192);
-    *reinterpret_cast<half2*>(emb_data) =
-        llvm_amdgcn_raw_buffer_load_fp16x2(emb_res, lane_id * sizeof(half2));
-    emb_data[2] = llvm_amdgcn_raw_buffer_load_fp16(
-        emb_res, (lane_id + 128) * sizeof(half));
-  }
-};
-
-template <typename index_t>
-struct load_row_per_warp<half, 256, index_t> {
-  static __device__ void
-  run(half* emb_data, index_t row_index, const half* p_emb_table, int lane_id) {
-    int32x4_t emb_res =
-        amdgcn_make_buffer_resource(p_emb_table + row_index * 256);
-    *reinterpret_cast<half2*>(&emb_data[0]) =
-        llvm_amdgcn_raw_buffer_load_fp16x2(emb_res, lane_id * sizeof(half2));
-    *reinterpret_cast<half2*>(&emb_data[2]) =
-        llvm_amdgcn_raw_buffer_load_fp16x2(
-            emb_res, (lane_id + 64) * sizeof(half2));
-  }
-};
-
-template <typename index_t>
-struct load_row_per_warp<half, 320, index_t> {
+template <int32_t embedding_dim, typename index_t>
+struct load_row_per_warp<half, embedding_dim, index_t> {
+  static constexpr int dword_per_row =
+      (embedding_dim + THREADS_PER_ROW - 1) / THREADS_PER_ROW;
+  static constexpr int half2_pairs_per_row = dword_per_row / 2;
+  static constexpr bool has_scalar_tail = dword_per_row % 2 == 1;
   static __device__ void
   run(half* emb_data, index_t row_index, const half* p_emb_table, int lane_id) {
     int32x4_t emb_res = amdgcn_make_buffer_resource(
-        p_emb_table + row_index * 320, sizeof(half) * 320);
-    *reinterpret_cast<half2*>(&emb_data[0]) =
-        llvm_amdgcn_raw_buffer_load_fp16x2(emb_res, lane_id * sizeof(half2));
-    *reinterpret_cast<half2*>(&emb_data[2]) =
-        llvm_amdgcn_raw_buffer_load_fp16x2(
-            emb_res, (lane_id + 64) * sizeof(half2));
-    emb_data[4] = llvm_amdgcn_raw_buffer_load_fp16(
-        emb_res, (lane_id + 256) * sizeof(half));
+        p_emb_table + row_index * embedding_dim, sizeof(half) * embedding_dim);
+#pragma unroll
+    for (int p = 0; p < half2_pairs_per_row; p++) {
+      *reinterpret_cast<half2*>(&emb_data[2 * p]) =
+          llvm_amdgcn_raw_buffer_load_fp16x2(
+              emb_res, (lane_id + p * 64) * sizeof(half2));
+    }
+    if constexpr (has_scalar_tail) {
+      constexpr int tail = dword_per_row - 1;
+      emb_data[tail] = llvm_amdgcn_raw_buffer_load_fp16(
+          emb_res, (lane_id + tail * 64) * sizeof(half));
+    }
   }
 };
 
@@ -253,111 +205,45 @@ struct load_row_per_warp<c10::Half, embedding_dim, index_t> {
   }
 };
 
-template <typename index_t>
-struct load_row_per_warp<float, 64, index_t> {
-  static __device__ void run(
-      float* emb_data,
-      index_t row_index,
-      const float* p_emb_table,
-      int lane_id) {
-    int32x4_t emb_res =
-        amdgcn_make_buffer_resource(p_emb_table + row_index * 64);
-    emb_data[0] =
-        llvm_amdgcn_raw_buffer_load_fp32(emb_res, lane_id * sizeof(float));
-  }
-};
-
-template <typename index_t>
-struct load_row_per_warp<float, 128, index_t> {
-  static __device__ void run(
-      float* emb_data,
-      index_t row_index,
-      const float* p_emb_table,
-      int lane_id) {
-    int32x4_t emb_res =
-        amdgcn_make_buffer_resource(p_emb_table + row_index * 128);
-    emb_data[0] =
-        llvm_amdgcn_raw_buffer_load_fp32(emb_res, lane_id * sizeof(float));
-    emb_data[1] = llvm_amdgcn_raw_buffer_load_fp32(
-        emb_res, (lane_id + 64) * sizeof(float));
-  }
-};
-
-template <typename index_t>
-struct load_row_per_warp<float, 160, index_t> {
-  static __device__ void run(
-      float* emb_data,
-      index_t row_index,
-      const float* p_emb_table,
-      int lane_id) {
-    int32x4_t emb_res = amdgcn_make_buffer_resource(
-        p_emb_table + row_index * 160, sizeof(float) * 160);
-    emb_data[0] =
-        llvm_amdgcn_raw_buffer_load_fp32(emb_res, lane_id * sizeof(float));
-    emb_data[1] = llvm_amdgcn_raw_buffer_load_fp32(
-        emb_res, (lane_id + 64) * sizeof(float));
-    emb_data[2] = llvm_amdgcn_raw_buffer_load_fp32(
-        emb_res, (lane_id + 128) * sizeof(float));
-  }
-};
-
-template <typename index_t>
-struct load_row_per_warp<float, 192, index_t> {
-  static __device__ void run(
-      float* emb_data,
-      index_t row_index,
-      const float* p_emb_table,
-      int lane_id) {
-    int32x4_t emb_res =
-        amdgcn_make_buffer_resource(p_emb_table + row_index * 192);
-    emb_data[0] =
-        llvm_amdgcn_raw_buffer_load_fp32(emb_res, lane_id * sizeof(float));
-    emb_data[1] = llvm_amdgcn_raw_buffer_load_fp32(
-        emb_res, (lane_id + 64) * sizeof(float));
-    emb_data[2] = llvm_amdgcn_raw_buffer_load_fp32(
-        emb_res, (lane_id + 128) * sizeof(float));
-  }
-};
-
-template <typename index_t>
-struct load_row_per_warp<float, 256, index_t> {
+template <int32_t embedding_dim, typename index_t>
+struct load_row_per_warp<c10::BFloat16, embedding_dim, index_t> {
   static __device__ void run(
-      float* emb_data,
+      c10::BFloat16* emb_data,
       index_t row_index,
-      const float* p_emb_table,
+      const c10::BFloat16* p_emb_table,
       int lane_id) {
-    int32x4_t emb_res =
-        amdgcn_make_buffer_resource(p_emb_table + row_index * 256);
-    emb_data[0] =
-        llvm_amdgcn_raw_buffer_load_fp32(emb_res, lane_id * sizeof(float));
-    emb_data[1] = llvm_amdgcn_raw_buffer_load_fp32(
-        emb_res, (lane_id + 64) * sizeof(float));
-    emb_data[2] = llvm_amdgcn_raw_buffer_load_fp32(
-        emb_res, (lane_id + 128) * sizeof(float));
-    emb_data[3] = llvm_amdgcn_raw_buffer_load_fp32(
-        emb_res, (lane_id + 192) * sizeof(float));
+    load_row_per_warp<half, embedding_dim, index_t>::run(
+        reinterpret_cast<half*>(emb_data),
+        row_index,
+        reinterpret_cast<const half*>(p_emb_table),
+        lane_id);
   }
 };
 
-template <typename index_t>
-struct load_row_per_warp<float, 320, index_t> {
+template <int32_t embedding_dim, typename index_t>
+struct load_row_per_warp<float, embedding_dim, index_t> {
+  static constexpr int dword_per_row =
+      (embedding_dim + THREADS_PER_ROW - 1) / THREADS_PER_ROW;
+  static constexpr int float2_pairs_per_row = dword_per_row / 2;
+  static constexpr bool has_scalar_tail = dword_per_row % 2 == 1;
   static __device__ void run(
       float* emb_data,
       index_t row_index,
       const float* p_emb_table,
       int lane_id) {
     int32x4_t emb_res = amdgcn_make_buffer_resource(
-        p_emb_table + row_index * 320, sizeof(float) * 320);
-    emb_data[0] =
-        llvm_amdgcn_raw_buffer_load_fp32(emb_res, lane_id * sizeof(float));
-    emb_data[1] = llvm_amdgcn_raw_buffer_load_fp32(
-        emb_res, (lane_id + 64) * sizeof(float));
-    emb_data[2] = llvm_amdgcn_raw_buffer_load_fp32(
-        emb_res, (lane_id + 128) * sizeof(float));
-    emb_data[3] = llvm_amdgcn_raw_buffer_load_fp32(
-        emb_res, (lane_id + 192) * sizeof(float));
-    emb_data[4] = llvm_amdgcn_raw_buffer_load_fp32(
-        emb_res, (lane_id + 256) * sizeof(float));
+        p_emb_table + row_index * embedding_dim, sizeof(float) * embedding_dim);
+#pragma unroll
+    for (int p = 0; p < float2_pairs_per_row; p++) {
+      *reinterpret_cast<floatx2_t*>(&emb_data[2 * p]) =
+          llvm_amdgcn_raw_buffer_load_fp32x2(
+              emb_res, (lane_id + p * 64) * sizeof(floatx2_t));
+    }
+    if constexpr (has_scalar_tail) {
+      constexpr int tail = dword_per_row - 1;
+      emb_data[tail] = llvm_amdgcn_raw_buffer_load_fp32(
+          emb_res, (lane_id + tail * 64) * sizeof(float));
+    }
   }
 };
 
@@ -408,72 +294,27 @@ struct store_row_per_warp {
   }
 };
 
-template <>
-struct store_row_per_warp<half, 64> {
-  static __device__ void run(const half* acc, half* p_output, int lane_id) {
-    int32x4_t out_res = amdgcn_make_buffer_resource(p_output);
-    llvm_amdgcn_raw_buffer_store_fp16(acc[0], out_res, lane_id * sizeof(half));
-  }
-};
-
-template <>
-struct store_row_per_warp<half, 128> {
-  static __device__ void run(const half* acc, half* p_output, int lane_id) {
-    int32x4_t out_res = amdgcn_make_buffer_resource(p_output);
-    llvm_amdgcn_raw_buffer_store_fp16x2(
-        *reinterpret_cast<const half2*>(acc), out_res, lane_id * sizeof(half2));
-  }
-};
-
-template <>
-struct store_row_per_warp<half, 160> {
-  static __device__ void run(const half* acc, half* p_output, int lane_id) {
-    int32x4_t out_res =
-        amdgcn_make_buffer_resource(p_output, 160 * sizeof(half));
-    llvm_amdgcn_raw_buffer_store_fp16x2(
-        *reinterpret_cast<const half2*>(acc), out_res, lane_id * sizeof(half2));
-    llvm_amdgcn_raw_buffer_store_fp16(
-        acc[2], out_res, (lane_id + 128) * sizeof(half));
-  }
-};
-
-template <>
-struct store_row_per_warp<half, 192> {
-  static __device__ void run(const half* acc, half* p_output, int lane_id) {
-    int32x4_t out_res = amdgcn_make_buffer_resource(p_output);
-    llvm_amdgcn_raw_buffer_store_fp16x2(
-        *reinterpret_cast<const half2*>(acc), out_res, lane_id * sizeof(half2));
-    llvm_amdgcn_raw_buffer_store_fp16(
-        acc[2], out_res, (lane_id + 128) * sizeof(half));
-  }
-};
-
-template <>
-struct store_row_per_warp<half, 256> {
-  static __device__ void run(const half* acc, half* p_output, int lane_id) {
-    int32x4_t out_res = amdgcn_make_buffer_resource(p_output);
-    llvm_amdgcn_raw_buffer_store_fp16x2(
-        *reinterpret_cast<const half2*>(acc), out_res, lane_id * sizeof(half2));
-    llvm_amdgcn_raw_buffer_store_fp16x2(
-        *reinterpret_cast<const half2*>(acc + 2),
-        out_res,
-        (lane_id + 64) * sizeof(half2));
-  }
-};
-
-template <>
-struct store_row_per_warp<half, 320> {
+template <int32_t embedding_dim>
+struct store_row_per_warp<half, embedding_dim> {
+  static constexpr int dword_per_row =
+      (embedding_dim + THREADS_PER_ROW - 1) / THREADS_PER_ROW;
+  static constexpr int half2_pairs_per_row = dword_per_row / 2;
+  static constexpr bool has_scalar_tail = dword_per_row % 2 == 1;
   static __device__ void run(const half* acc, half* p_output, int lane_id) {
     int32x4_t out_res =
-        amdgcn_make_buffer_resource(p_output, 320 * sizeof(half));
-    llvm_amdgcn_raw_buffer_store_fp16x2(
-        *reinterpret_cast<const half2*>(acc), out_res, lane_id * sizeof(half2));
-    llvm_amdgcn_raw_buffer_store_fp16x2(
-        *reinterpret_cast<const half2*>(acc + 2),
-        out_res,
-        (lane_id + 64) * sizeof(half2));
-    llvm_amdgcn_raw_buffer_store_fp16(
-        acc[4], out_res, (lane_id + 256) * sizeof(half));
+        amdgcn_make_buffer_resource(p_output, sizeof(half) * embedding_dim);
+#pragma unroll
+    for (int p = 0; p < half2_pairs_per_row; p++) {
+      llvm_amdgcn_raw_buffer_store_fp16x2(
+          *reinterpret_cast<const half2*>(acc + 2 * p),
+          out_res,
+          (lane_id + p * 64) * sizeof(half2));
+    }
+    if constexpr (has_scalar_tail) {
+      constexpr int tail = dword_per_row - 1;
+      llvm_amdgcn_raw_buffer_store_fp16(
+          acc[tail], out_res, (lane_id + tail * 64) * sizeof(half));
+    }
   }
 };
 
@@ -488,77 +329,27 @@ struct store_row_per_warp<c10::Half, embedding_dim> {
   }
 };
 
-template <>
-struct store_row_per_warp<float, 64> {
-  static __device__ void run(const float* acc, float* p_output, int lane_id) {
-    int32x4_t out_res = amdgcn_make_buffer_resource(p_output);
-    llvm_amdgcn_raw_buffer_store_fp32(acc[0], out_res, lane_id * sizeof(float));
-  }
-};
-
-template <>
-struct store_row_per_warp<float, 128> {
-  static __device__ void run(const float* acc, float* p_output, int lane_id) {
-    int32x4_t out_res = amdgcn_make_buffer_resource(p_output);
-    llvm_amdgcn_raw_buffer_store_fp32(acc[0], out_res, lane_id * sizeof(float));
-    llvm_amdgcn_raw_buffer_store_fp32(
-        acc[1], out_res, (lane_id + 64) * sizeof(float));
-  }
-};
-
-template <>
-struct store_row_per_warp<float, 160> {
-  static __device__ void run(const float* acc, float* p_output, int lane_id) {
-    int32x4_t out_res =
-        amdgcn_make_buffer_resource(p_output, sizeof(float) * 160);
-    llvm_amdgcn_raw_buffer_store_fp32(acc[0], out_res, lane_id * sizeof(float));
-    llvm_amdgcn_raw_buffer_store_fp32(
-        acc[1], out_res, (lane_id + 64) * sizeof(float));
-    llvm_amdgcn_raw_buffer_store_fp32(
-        acc[2], out_res, (lane_id + 128) * sizeof(float));
-  }
-};
-
-template <>
-struct store_row_per_warp<float, 192> {
-  static __device__ void run(const float* acc, float* p_output, int lane_id) {
-    int32x4_t out_res = amdgcn_make_buffer_resource(p_output);
-    llvm_amdgcn_raw_buffer_store_fp32(acc[0], out_res, lane_id * sizeof(float));
-    llvm_amdgcn_raw_buffer_store_fp32(
-        acc[1], out_res, (lane_id + 64) * sizeof(float));
-    llvm_amdgcn_raw_buffer_store_fp32(
-        acc[2], out_res, (lane_id + 128) * sizeof(float));
-  }
-};
-
-template <>
-struct store_row_per_warp<float, 256> {
-  static __device__ void run(const float* acc, float* p_output, int lane_id) {
-    int32x4_t out_res = amdgcn_make_buffer_resource(p_output);
-    llvm_amdgcn_raw_buffer_store_fp32(acc[0], out_res, lane_id * sizeof(float));
-    llvm_amdgcn_raw_buffer_store_fp32(
-        acc[1], out_res, (lane_id + 64) * sizeof(float));
-    llvm_amdgcn_raw_buffer_store_fp32(
-        acc[2], out_res, (lane_id + 128) * sizeof(float));
-    llvm_amdgcn_raw_buffer_store_fp32(
-        acc[3], out_res, (lane_id + 192) * sizeof(float));
-  }
-};
-
-template <>
-struct store_row_per_warp<float, 320> {
+template <int32_t embedding_dim>
+struct store_row_per_warp<float, embedding_dim> {
+  static constexpr int dword_per_row =
+      (embedding_dim + THREADS_PER_ROW - 1) / THREADS_PER_ROW;
+  static constexpr int float2_pairs_per_row = dword_per_row / 2;
+  static constexpr bool has_scalar_tail = dword_per_row % 2 == 1;
   static __device__ void run(const float* acc, float* p_output, int lane_id) {
     int32x4_t out_res =
-        amdgcn_make_buffer_resource(p_output, sizeof(float) * 320);
-    llvm_amdgcn_raw_buffer_store_fp32(acc[0], out_res, lane_id * sizeof(float));
-    llvm_amdgcn_raw_buffer_store_fp32(
-        acc[1], out_res, (lane_id + 64) * sizeof(float));
-    llvm_amdgcn_raw_buffer_store_fp32(
-        acc[2], out_res, (lane_id + 128) * sizeof(float));
-    llvm_amdgcn_raw_buffer_store_fp32(
-        acc[3], out_res, (lane_id + 192) * sizeof(float));
-    llvm_amdgcn_raw_buffer_store_fp32(
-        acc[4], out_res, (lane_id + 256) * sizeof(float));
+        amdgcn_make_buffer_resource(p_output, sizeof(float) * embedding_dim);
+#pragma unroll
+    for (int p = 0; p < float2_pairs_per_row; p++) {
+      llvm_amdgcn_raw_buffer_store_fp32x2(
+          *reinterpret_cast<const floatx2_t*>(acc + 2 * p),
+          out_res,
+          (lane_id + p * 64) * sizeof(floatx2_t));
+    }
+    if constexpr (has_scalar_tail) {
+      constexpr int tail = dword_per_row - 1;
+      llvm_amdgcn_raw_buffer_store_fp32(
+          acc[tail], out_res, (lane_id + tail * 64) * sizeof(float));
+    }
   }
 };
 
diff --git a/fbgemm_gpu/test/tbe/training/backward_adagrad_test.py b/fbgemm_gpu/test/tbe/training/backward_adagrad_test.py
index 1e8f5e87e5..a4701dae19 100644
--- a/fbgemm_gpu/test/tbe/training/backward_adagrad_test.py
+++ b/fbgemm_gpu/test/tbe/training/backward_adagrad_test.py
@@ -13,7 +13,7 @@
 import os
 import time
 import unittest
-from typing import Any
+from typing import Any, Optional
 
 import torch
 from fbgemm_gpu.split_embedding_configs import (
@@ -247,6 +247,7 @@ def _test_backward_adagrad_rocm_kernel(
         weights_precision: SparseType,
         weighted: bool,
         weight_decay_mode: WeightDecayMode,
+        output_dtype: Optional[SparseType] = None,
     ) -> None:
         """Helper method for ROCm backward kernel tests."""
         execute_backward_adagrad(
@@ -266,7 +267,9 @@ def _test_backward_adagrad_rocm_kernel(
             cache_algorithm=CacheAlgorithm.LRU,
             pooling_mode=PoolingMode.SUM,
             use_cpu=False,
-            output_dtype=weights_precision,
+            output_dtype=(
+                output_dtype if output_dtype is not None else weights_precision
+            ),
             weight_decay_mode=weight_decay_mode,
         )
 
@@ -277,6 +280,9 @@ def _test_backward_adagrad_rocm_kernel(
         log_E=st.integers(min_value=3, max_value=5),
         L=st.integers(min_value=2, max_value=20),
         weights_precision=st.sampled_from([SparseType.FP16, SparseType.FP32]),
+        output_dtype=st.sampled_from(
+            [SparseType.FP32, SparseType.FP16, SparseType.BF16]
+        ),
         weighted=st.booleans(),
         weight_decay_mode=st.sampled_from(
             [
@@ -297,6 +303,7 @@ def test_backward_adagrad_rocm_fallback_kernel(
         log_E: int,
         L: int,
         weights_precision: SparseType,
+        output_dtype: SparseType,
         weighted: bool,
         weight_decay_mode: WeightDecayMode,
     ) -> None:
@@ -315,6 +322,7 @@ def test_backward_adagrad_rocm_fallback_kernel(
                 weights_precision=weights_precision,
                 weighted=weighted,
                 weight_decay_mode=weight_decay_mode,
+                output_dtype=output_dtype,
             )
 
     @given(
@@ -324,6 +332,9 @@ def test_backward_adagrad_rocm_fallback_kernel(
         log_E=st.integers(min_value=3, max_value=5),
         L=st.integers(min_value=2, max_value=20),
         weights_precision=st.sampled_from([SparseType.FP16, SparseType.FP32]),
+        output_dtype=st.sampled_from(
+            [SparseType.FP32, SparseType.FP16, SparseType.BF16]
+        ),
         weighted=st.booleans(),
         weight_decay_mode=st.sampled_from(
             [
@@ -344,6 +355,7 @@ def test_backward_adagrad_rocm_optimized_kernel(
         log_E: int,
         L: int,
         weights_precision: SparseType,
+        output_dtype: SparseType,
         weighted: bool,
         weight_decay_mode: WeightDecayMode,
     ) -> None:
@@ -362,6 +374,7 @@ def test_backward_adagrad_rocm_optimized_kernel(
                 weights_precision=weights_precision,
                 weighted=weighted,
                 weight_decay_mode=weight_decay_mode,
+                output_dtype=output_dtype,
             )
 
     @given(