move common op to vector utils (pytorch#3759)

Summary:
Pull Request resolved: pytorch#3759

X-link: facebookresearch/FBGEMM#840

as title. Op is used in next diff.

Reviewed By: sgrigory

Differential Revision: D70387235

fbshipit-source-id: c5acebd1e371e67822d51f3b5cb6f194c99ef1eb

Author: Aya-ZIbra

fbgemm_gpu/experimental/gen_ai/src/kv_cache/kv_cache.cu

@@ -546,13 +546,7 @@ DEVICE_INLINE fx4 rope_xpos(
     double hi_freq_factor = 32) {
   fx4 dst; // read 4 bf16 from src and store in 4 float registers
   if (head == QKV::V) {
-    auto r0 = bf1622float2(src.vals[0]);
-    auto r1 = bf1622float2(src.vals[1]);
-    dst.x = r0.x;
-    dst.y = r0.y;
-    dst.z = r1.x;
-    dst.w = r1.y;
-    return dst;
+    return bfx4_to_fx4(src);
   }
   int32_t offset_0 = ((4 * threadIdx.x) / 2 + 0);
   int32_t offset_1 = ((4 * threadIdx.x) / 2 + 1);

fbgemm_gpu/include/fbgemm_gpu/utils/vec_quant.cuh

@@ -173,6 +173,16 @@ DEVICE_INLINE fx4 bfx4_scale_acc(fx4 acc, bfx4 a, float b) {
   return acc;
 }
 
+DEVICE_INLINE fx4 bfx4_to_fx4(bfx4 src) {
+  fx4 dst;
+  auto r0 = bf1622float2(src.vals[0]);
+  auto r1 = bf1622float2(src.vals[1]);
+  dst.x = r0.x;
+  dst.y = r0.y;
+  dst.z = r1.x;
+  dst.w = r1.y;
+  return dst;
+}
 DEVICE_INLINE fx4 fx4_acc(fx4 a, fx4 b) {
   a.x += b.x;
   a.y += b.y;