merge redundant files

Signed-off-by: Zhu, Zufang <[email protected]>

Author: zufangzhu

csrc/xpu/cache.cpp

@@ -5,7 +5,7 @@
 #include <string>
 
 #include "dispatch_utils.h"
-#include "quantization/fp8/quant_utils.hpp"
+#include "quantization/fp8/quant_utils.h"
 #include "utils.h"
 
 namespace vllm {
@@ -90,6 +90,25 @@ void call_reshape_and_cache(
   });
 }
 
+// Used by vectorization_utils to copy/convert one element
+template <typename OutT, typename InT, Fp8KVCacheDataType kv_dt>
+struct CopyWithScaleOp {
+  float scale;
+
+  inline void operator()(OutT& dst, const InT src) const {
+    if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
+      dst = static_cast<OutT>(src);
+    } else {
+      float x = (float)src / scale;
+      if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) {
+        dst = static_cast<at::Float8_e4m3fn>(x);
+      } else if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E5M2) {
+        dst = static_cast<at::Float8_e5m2>(x);
+      }
+    }
+  }
+};
+
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
 void reshape_and_cache_flash_kernel(
     const scalar_t* __restrict__ key,    // [num_tokens, num_heads, head_size]
@@ -127,8 +146,8 @@ void reshape_and_cache_flash_kernel(
   float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale;
   float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *v_scale;
 
-  fp8::CopyWithScaleOp<cache_t, scalar_t, kv_dt> k_op{k_scale_val};
-  fp8::CopyWithScaleOp<cache_t, scalar_t, kv_dt> v_op{v_scale_val};
+  CopyWithScaleOp<cache_t, scalar_t, kv_dt> k_op{k_scale_val};
+  CopyWithScaleOp<cache_t, scalar_t, kv_dt> v_op{v_scale_val};
   fp8::scaled_convert_vec(key_src, key_dst, n, local_idx, local_range, k_op);
   fp8::scaled_convert_vec(value_src, value_dst, n, local_idx, local_range,
                           v_op);

csrc/xpu/dispatch_utils.h

@@ -21,13 +21,13 @@
 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
   AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
 
-#define VLLM_DISPATCH_CASE_FP8_TYPES(...) \
+#define VLLM_DISPATCH_CASE_FP8_TYPES(...)                          \
   AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
   AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e5m2, __VA_ARGS__)
 
 #define VLLM_DISPATCH_CASE_QUANT_TYPES(...)                    \
   AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::Float8_e5m2, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::Float8_e5m2, __VA_ARGS__)   \
   AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
 
 // When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'.

csrc/xpu/quantization/fp8/fp8_quant.cpp

@@ -8,8 +8,7 @@
 #include "xpu/ops.h"
 
 #include "fp8_quant.h"
-#include "utils.h"
-
+#include "quant_utils.h"
 
 namespace vllm {
 
@@ -22,25 +21,19 @@ class scaled_fp8_quant_kernel {
   int64_t num_elems;
 
  public:
-  scaled_fp8_quant_kernel(
-      fp8_type* out_,
-      const scalar_t* input_,
-      const float* scale_,
-      int64_t num_elems_)
+  scaled_fp8_quant_kernel(fp8_type* out_, const scalar_t* input_,
+                          const float* scale_, int64_t num_elems_)
       : out(out_), input(input_), scale(scale_), num_elems(num_elems_) {}
   void operator()(sycl::nd_item<1> item) const {
     int tid = item.get_global_linear_id();
 
     // Invert the scale so that we can use multiplications to avoid expensive
     // division.
     const float inverted_scale = 1.0f / (*scale);
-    scaled_fp8_conversion_vec<scalar_t, true>(
-        out,
-        input,
-        inverted_scale,
-        num_elems,
-        tid,
-        item.get_local_range(0) * item.get_group_range(0));
+    fp8::ConvertWithScaleOp<true, fp8_type> op{inverted_scale};
+    fp8::scaled_convert_vec(input, out, num_elems, tid,
+                            item.get_local_range(0) * item.get_group_range(0),
+                            op);
   }
 };
 
@@ -54,12 +47,10 @@ class dynamic_per_token_scaled_fp8_quant_kernel {
   const int hidden_size;
 
  public:
-  dynamic_per_token_scaled_fp8_quant_kernel(
-      fp8_type* out_,
-      float* scale_,
-      scalar_t const* input_,
-      float const* scale_ub_,
-      const int hidden_size_)
+  dynamic_per_token_scaled_fp8_quant_kernel(fp8_type* out_, float* scale_,
+                                            scalar_t const* input_,
+                                            float const* scale_ub_,
+                                            const int hidden_size_)
       : out(out_),
         scale(scale_),
         input(input_),
@@ -70,13 +61,6 @@ class dynamic_per_token_scaled_fp8_quant_kernel {
     int const tid = item.get_local_id(0);
     int const token_idx = item.get_group(0);
 
-    // sycl::ext::oneapi::experimental::printf(
-    //     "token_idx: %d, tid: %d, hidden_size: %d, group_range: %d\n",
-    //     token_idx,
-    //     tid,
-    //     hidden_size,
-    //     item.get_local_range(0));
-
     // Use int64 to avoid overflowing an int32 when calculating this offset
     int64_t offset = static_cast<int64_t>(token_idx) * hidden_size;
     scalar_t const* token_input = &input[offset];
@@ -88,8 +72,8 @@ class dynamic_per_token_scaled_fp8_quant_kernel {
 
     float absmax_val = 0.0f;
     if (can_vectorize) {
-      absmax_val = thread_max_vec(
-          token_input, hidden_size, tid, item.get_local_range(0));
+      absmax_val = thread_max_vec(token_input, hidden_size, tid,
+                                  item.get_local_range(0));
     } else {
       for (int i = tid; i < hidden_size; i += item.get_local_range(0)) {
         float const x = static_cast<float>(token_input[i]);
@@ -110,38 +94,33 @@ class dynamic_per_token_scaled_fp8_quant_kernel {
         token_scale[0] = block_absmax_val_maybe;
       }
       // token scale computation
-      token_scale[0] = sycl::max(
-          token_scale[0] / quant_type_max_v<fp8_type>,
-          min_scaling_factor<fp8_type>::val());
+      token_scale[0] =
+          sycl::max(token_scale[0] / fp8::quant_type_max_v<fp8_type>,
+                    fp8::min_scaling_factor<fp8_type>::val());
       scale[token_idx] = token_scale[0];
     }
     group_barrier(item.get_group());
 
     // Note that we don't use inverted scales so we can match FBGemm impl.
     const float inverted_scale = 1.0f / (token_scale[0]);
     if (can_vectorize) {
-      scaled_fp8_conversion_vec<scalar_t, true>(
-          token_output,
-          token_input,
-          inverted_scale,
-          hidden_size,
-          tid,
-          item.get_local_range(0));
+      fp8::ConvertWithScaleOp<true, fp8_type> op{inverted_scale};
+      fp8::scaled_convert_vec(token_input, token_output, hidden_size, tid,
+                              item.get_local_range(0), op);
     } else {
       for (int i = tid; i < hidden_size; i += item.get_local_range(0)) {
-        token_output[i] = scaled_fp8_conversion<true, fp8_type>(
-            static_cast<float>(token_input[i]), inverted_scale, tid, token_idx);
+        fp8::ConvertWithScaleOp<true, fp8_type> op{inverted_scale};
+        op(token_output[i], token_input[i]);
       }
     }
   }
 };
 
-} // namespace vllm
+}  // namespace vllm
 
-void static_scaled_fp8_quant(
-    torch::Tensor& out, // [..., d]
-    torch::Tensor const& input, // [..., d]
-    torch::Tensor const& scale) // [1]
+void static_scaled_fp8_quant(torch::Tensor& out,          // [..., d]
+                             torch::Tensor const& input,  // [..., d]
+                             torch::Tensor const& scale)  // [1]
 {
   int64_t num_tokens = input.numel() / input.size(-1);
   int64_t num_elems = input.numel();
@@ -158,21 +137,18 @@ void static_scaled_fp8_quant(
               // Launch the kernel
               stream.submit([&](sycl::handler& cgh) {
                 auto kernel = vllm::scaled_fp8_quant_kernel<scalar_t, fp8_t>(
-                    out.data_ptr<fp8_t>(),
-                    input.data_ptr<scalar_t>(),
-                    scale.data_ptr<float>(),
-                    num_elems);
-                cgh.parallel_for(
-                    sycl::nd_range<1>(grid * block, block), kernel);
+                    out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
+                    scale.data_ptr<float>(), num_elems);
+                cgh.parallel_for(sycl::nd_range<1>(grid * block, block),
+                                 kernel);
               });
             });
       });
 }
 
-void dynamic_scaled_fp8_quant(
-    torch::Tensor& out, // [..., d]
-    torch::Tensor const& input, // [..., d]
-    torch::Tensor& scale) // [1]
+void dynamic_scaled_fp8_quant(torch::Tensor& out,          // [..., d]
+                              torch::Tensor const& input,  // [..., d]
+                              torch::Tensor& scale)        // [1]
 {
   int64_t num_tokens = input.numel() / input.size(-1);
   int64_t num_elems = input.numel();
@@ -190,30 +166,26 @@ void dynamic_scaled_fp8_quant(
               stream.submit([&](sycl::handler& cgh) {
                 auto max_reduce_kernel =
                     vllm::segmented_max_reduction<scalar_t, fp8_t>(
-                        scale.data_ptr<float>(),
-                        input.data_ptr<scalar_t>(),
+                        scale.data_ptr<float>(), input.data_ptr<scalar_t>(),
                         num_elems);
-                cgh.parallel_for(
-                    sycl::nd_range<1>(grid * block, block), max_reduce_kernel);
+                cgh.parallel_for(sycl::nd_range<1>(grid * block, block),
+                                 max_reduce_kernel);
               });
               stream.submit([&](sycl::handler& cgh) {
                 auto kernel = vllm::scaled_fp8_quant_kernel<scalar_t, fp8_t>(
-                    out.data_ptr<fp8_t>(),
-                    input.data_ptr<scalar_t>(),
-                    scale.data_ptr<float>(),
-                    num_elems);
-                cgh.parallel_for(
-                    sycl::nd_range<1>(grid * block, block), kernel);
+                    out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
+                    scale.data_ptr<float>(), num_elems);
+                cgh.parallel_for(sycl::nd_range<1>(grid * block, block),
+                                 kernel);
               });
             });
       });
 }
 
 void dynamic_per_token_scaled_fp8_quant(
-    torch::Tensor& out, // [..., d]
-    torch::Tensor const& input, // [..., d]
-    torch::Tensor& scales,
-    std::optional<at::Tensor> const& scale_ub) {
+    torch::Tensor& out,          // [..., d]
+    torch::Tensor const& input,  // [..., d]
+    torch::Tensor& scales, std::optional<at::Tensor> const& scale_ub) {
   TORCH_CHECK(input.is_contiguous());
   TORCH_CHECK(out.is_contiguous());
 
@@ -228,26 +200,23 @@ void dynamic_per_token_scaled_fp8_quant(
   auto stream = at::xpu::getCurrentXPUStream().queue();
   VLLM_DISPATCH_FLOATING_TYPES(
       input.scalar_type(),
-      "dynamic_per_token_scaled_fp8_quant_kernel_scalar_type",
-      [&] {
+      "dynamic_per_token_scaled_fp8_quant_kernel_scalar_type", [&] {
         VLLM_DISPATCH_FP8_TYPES(
             out.scalar_type(),
-            "dynamic_per_token_scaled_fp8_quant_kernel_fp8_type",
-            [&] {
+            "dynamic_per_token_scaled_fp8_quant_kernel_fp8_type", [&] {
               // Launch the kernel
               stream
                   .submit([&](sycl::handler& cgh) {
-                    auto kernel = vllm::dynamic_per_token_scaled_fp8_quant_kernel<
-                        scalar_t,
-                        fp8_t>(
-                        out.data_ptr<fp8_t>(),
-                        scales.data_ptr<float>(),
-                        input.data_ptr<scalar_t>(),
-                        scale_ub.has_value() ? scale_ub->data_ptr<float>()
-                                             : nullptr,
-                        hidden_size);
-                    cgh.parallel_for(
-                        sycl::nd_range<1>(grid * block, block), kernel);
+                    auto kernel =
+                        vllm::dynamic_per_token_scaled_fp8_quant_kernel<
+                            scalar_t, fp8_t>(
+                            out.data_ptr<fp8_t>(), scales.data_ptr<float>(),
+                            input.data_ptr<scalar_t>(),
+                            scale_ub.has_value() ? scale_ub->data_ptr<float>()
+                                                 : nullptr,
+                            hidden_size);
+                    cgh.parallel_for(sycl::nd_range<1>(grid * block, block),
+                                     kernel);
                   })
                   .wait();
             });