grouped topk from IPEX (#33)

* add total fused grouped topk

Signed-off-by: mayuyuace <[email protected]>

* fix warnings

Signed-off-by: mayuyuace <[email protected]>

* format

Signed-off-by: mayuyuace <[email protected]>

---------

Signed-off-by: mayuyuace <[email protected]>

Author: mayuyuace

CMakeLists.txt

@@ -222,6 +222,7 @@ define_gpu_extension_target(
 set(VLLM_MOE_EXT_SRC
   "csrc/moe/torch_bindings.cpp"
   "csrc/moe/grouped_topk.cpp"
+  "csrc/moe/fused_grouped_topk.cpp"
   "csrc/moe/moe_align_sum_kernels.cpp")
 
 message(STATUS "Enabling moe extension.")

benchmark/benchmark_grouped_topk.py

@@ -8,7 +8,8 @@
 import torch
 import triton
 
-from tests.ops.grouped_topk_op import fused_grouped_topk, grouped_topk
+from tests.ops.grouped_topk_op import (fused_grouped_topk,
+                                       fused_grouped_topk_sycl, grouped_topk)
 
 
 @torch.compile
@@ -95,8 +96,8 @@ def get_benchmark():
             ],
             x_vals=[tuple(_) for _ in configs],
             line_arg="provider",
-            line_vals=["vllm", "native", "compile"],
-            line_names=["vllm", "native", "compile"],
+            line_vals=["vllm", "native", "compile", "sycl"],
+            line_names=["vllm", "native", "compile", "sycl"],
             styles=[("blue", "-"), ("green", "-"), ("orange", "-"),
                     ("red", "-")],
             ylabel="us",
@@ -156,7 +157,7 @@ def benchmark(
                     e_score_correction_bias=e_score_correction_bias),
                 quantiles=quantiles,
             )
-        else:
+        elif provider == "compile":
             ms, min_ms, max_ms = triton.testing.do_bench(
                 lambda: grouped_topk_compile(
                     hidden_states=hidden_states,
@@ -170,6 +171,20 @@ def benchmark(
                     e_score_correction_bias=e_score_correction_bias),
                 quantiles=quantiles,
             )
+        elif provider == "sycl":
+            ms, min_ms, max_ms = triton.testing.do_bench(
+                lambda: fused_grouped_topk_sycl(
+                    hidden_states=hidden_states,
+                    gating_output=gating_output,
+                    topk=topk,
+                    renormalize=renormalize,
+                    num_expert_group=num_expert_group,
+                    topk_group=topk_group,
+                    scoring_func=scoring_func,
+                    routed_scaling_factor=routed_scaling_factor,
+                    e_score_correction_bias=e_score_correction_bias),
+                quantiles=quantiles,
+            )
 
         return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 

csrc/activation.cpp

@@ -73,7 +73,7 @@ class act_kernel {
              const int d)
       : out_(out), input_(input), d_(d) {}
 
-  void operator() [[intel::reqd_sub_group_size(32)]] (
+  void operator() [[sycl::reqd_sub_group_size(32)]] (
       const sycl::nd_item<3>& item_ct1) const {
     const int64_t token_idx = item_ct1.get_group(2);
     for (int64_t idx = item_ct1.get_local_id(2); idx < d_;
@@ -98,7 +98,7 @@ class act_and_mul_kernel {
                      const int d)
       : out_(out), input_(input), d_(d) {}
 
-  void operator() [[intel::reqd_sub_group_size(32)]] (
+  void operator() [[sycl::reqd_sub_group_size(32)]] (
       const sycl::nd_item<3>& item_ct1) const {
     const int64_t token_idx = item_ct1.get_group(2);
     for (int64_t idx = item_ct1.get_local_id(2); idx < d_;

csrc/layernorm.cpp

@@ -23,9 +23,10 @@ class rms_norm_kernel {
         hidden_size(hidden_size_),
         s_variance(s_variance_) {}
 
-  void operator() [[intel::reqd_sub_group_size(32)]] (
+  void operator() [[sycl::reqd_sub_group_size(32)]] (
       const sycl::nd_item<3>& item_ct1) const {
-    float* s_variance_ptr = s_variance.get_pointer();
+    float* s_variance_ptr =
+        s_variance.template get_multi_ptr<sycl::access::decorated::no>().get();
     float variance = 0.0f;
 
     for (int idx = item_ct1.get_local_id(2); idx < hidden_size;
@@ -65,7 +66,7 @@ class rms_norm_kernel {
 template <typename scalar_t>
 void call_rms_norm_kernel(torch::Tensor& out, torch::Tensor& input,
                           torch::Tensor& weight, float epsilon) {
-  using sycl_t = vllm::xpu::SyclTypeTrait<scalar_t>::Type;
+  using sycl_t = typename vllm::xpu::SyclTypeTrait<scalar_t>::Type;
   int hidden_size = input.size(-1);
   int num_tokens = input.numel() / hidden_size;
   int64_t input_stride = input.stride(-2);
@@ -104,9 +105,10 @@ class fused_add_rms_norm_kernel {
         hidden_size(hidden_size_),
         s_variance(s_variance_) {}
 
-  void operator() [[intel::reqd_sub_group_size(32)]] (
+  void operator() [[sycl::reqd_sub_group_size(32)]] (
       const sycl::nd_item<3>& item_ct1) const {
-    float* s_variance_ptr = s_variance.get_pointer();
+    float* s_variance_ptr =
+        s_variance.template get_multi_ptr<sycl::access::decorated::no>().get();
     float variance = 0.0f;
 
     for (int idx = item_ct1.get_local_id(2); idx < hidden_size;
@@ -150,7 +152,7 @@ template <typename scalar_t>
 void call_fused_add_rms_norm_kernel(torch::Tensor& input,
                                     torch::Tensor& residual,
                                     torch::Tensor& weight, float epsilon) {
-  using sycl_t = vllm::xpu::SyclTypeTrait<scalar_t>::Type;
+  using sycl_t = typename vllm::xpu::SyclTypeTrait<scalar_t>::Type;
   int hidden_size = input.size(-1);
   int num_tokens = input.numel() / hidden_size;
   auto input_ptr = input.data_ptr<scalar_t>();