feat(lora) Add lora shrink & expand ops

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

Author: chaojun-zhang

CMakeLists.txt

@@ -183,6 +183,8 @@ define_gpu_extension_target(
 if(VLLM_GPU_LANG STREQUAL "SYCL")
   set(VLLM_EXT_XPU_SRC
     "csrc/xpu/torch_bindings.cpp"
+    "csrc/xpu/lora/lora_shrink.cpp"
+    "csrc/xpu/lora/lora_expand.cpp"
   )
   include_directories("/usr/include")
   set(CMPLR_ROOT $ENV{CMPLR_ROOT})
@@ -236,27 +238,3 @@ define_gpu_extension_target(
   INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
   USE_SABI 3
   WITH_SOABI)
-
-#
-# _lora_C extension
-#
-
-set(VLLM_LORA_EXT_SRC
-    "csrc/lora/torch_bindings.cpp"
-    "csrc/lora/lora_shrink.cpp"
-    "csrc/lora/lora_expand.cpp"
-)
-
-message(STATUS "Enabling lora extension.")
-define_gpu_extension_target(
-  _lora_C
-  DESTINATION vllm_xpu_kernels
-  LANGUAGE ${VLLM_GPU_LANG}
-  SOURCES ${VLLM_LORA_EXT_SRC}
-  COMPILE_FLAGS ${VLLM_GPU_FLAGS}
-  LINK_FLAGS ${VLLM_GPU_LINK_FLAGS}
-  ARCHITECTURES ${VLLM_GPU_ARCHES}
-  INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
-  INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
-  USE_SABI 3
-  WITH_SOABI)

benchmark/benchmark_lora.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
 import copy
@@ -17,7 +18,7 @@
 from utils import ArgPool, Bench, CudaGraphBenchParams
 from weight_shapes import WEIGHT_SHAPES
 
-from tests.register_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
+from tests.lora.lora_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
 
 DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())
 DEFAULT_TP_SIZES = [1]

csrc/lora/torch_bindings.cpp

@@ -166,16 +166,16 @@ class bgmv_expand_kernel {
     sycl::group_barrier(sg);
 
     if (vec_id == 0) {
-      accscalar_t result = 0;
+      float result = 0.0f;
 #pragma unroll
       for (uint32_t i = 0; i < workitem_per_hidden_; ++i) {
-        result += slm_[slm_base + i];
+        result += static_cast<float>(slm_[slm_base + i]);
       }
       const size_t out_off = static_cast<size_t>(batch_id) * output_hidden_ +
                              slice_offset_ + hidden_id;
       if (add_to_output_) {
-        outputs_[out_off] = static_cast<output_t>(
-            static_cast<accscalar_t>(outputs_[out_off]) + result);
+        result += static_cast<float>(outputs_[out_off]);
+        outputs_[out_off] = static_cast<output_t>(result);
       } else {
         outputs_[out_off] = static_cast<output_t>(result);
       }

csrc/lora/lora_expand.cpp renamed to csrc/xpu/lora/lora_expand.cpp

@@ -1,5 +1,6 @@
 #include "core/registration.h"
 #include "xpu/ops.h"
+#include "xpu/lora/lora_ops.h"
 
 #include <torch/library.h>
 #include <torch/version.h>
@@ -11,6 +12,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, xpu_ops) {
       "fp8_gemm_w8a16(Tensor! A, Tensor! B, bool trans_B, Tensor? B_scale_, "
       "Tensor? bias_) -> Tensor");
   xpu_ops.impl("fp8_gemm_w8a16", torch::kXPU, &fp8_gemm_w8a16);
+
+  xpu_ops.def(
+      "bgmv_shrink(Tensor! outputs, Tensor inputs, Tensor weights, Tensor "
+      "indices, float scale) -> ()");
+  xpu_ops.impl("bgmv_shrink", torch::kXPU, &bgmv_shrink);
+
+  xpu_ops.def(
+      "bgmv_expand(Tensor! outputs, Tensor inputs, Tensor weights, Tensor "
+      "indices, bool add_to_output) -> ()");
+  xpu_ops.impl("bgmv_expand", torch::kXPU, &bgmv_expand);
+
+  xpu_ops.def(
+      "bgmv_expand_slice(Tensor! outputs, Tensor inputs, Tensor weights, "
+      "Tensor indices, int slice_offset,bool add_to_output) -> ()");
+  xpu_ops.impl("bgmv_expand_slice", torch::kXPU, &bgmv_expand_slice);
 }
 
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

csrc/lora/lora_ops.h renamed to csrc/xpu/lora/lora_ops.h

@@ -272,7 +272,6 @@ def run(self):
 if _build_custom_ops():
     ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._C"))
     ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._moe_C"))
-    ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._lora_C"))
     ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._xpu_C"))
 
 if ext_modules:

csrc/lora/lora_shrink.cpp renamed to csrc/xpu/lora/lora_shrink.cpp

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+import vllm_xpu_kernels._xpu_C  # noqa: F401
+
+
+def bgmv_shrink(
+    inputs: torch.Tensor,
+    lora_a_weights: torch.Tensor,
+    output_tensor: torch.Tensor,
+    lora_indices_tensor: torch.Tensor,
+    scaling: float = 1.0,
+) -> None:
+    torch.ops._xpu_C.bgmv_shrink(
+        output_tensor,
+        inputs,
+        lora_a_weights,
+        lora_indices_tensor,
+        scaling,
+    )
+
+
+def bgmv_expand(
+    inputs: torch.Tensor,
+    lora_b_weights: torch.Tensor,
+    output_tensor: torch.Tensor,
+    lora_indices_tensor: torch.Tensor,
+    add_inputs: bool = True,
+):
+    """
+    Args:
+        inputs (torch.Tensor): Shape: `[batch_size, hidden_size]`.
+        lora_b_weights (torch.Tensor): Shape: `[lora_num, rank, hidden_size]`.
+        output_tensor (torch.Tensor): Shape: `[batch_size, rank]`.
+        lora_indices_tensor (torch.Tensor): Shape: `[batch_size]`.
+         The LoRA index corresponding to each batch. An index of -1 means
+            no lora should be applied.
+        add_inputs (bool, optional):  Defaults to False. adds the final lora
+            results to the output.
+
+    Semantics:
+      for i in range(inputs.size(0)):
+        output_tensor[i] =
+            inputs[i] @ lora_b_weights[lora_indices_tensor[i]]
+            + (inputs[i] if add_inputs else 0)
+    """
+    torch.ops._xpu_C.bgmv_expand(
+        output_tensor,
+        inputs,
+        lora_b_weights,
+        lora_indices_tensor,
+        add_inputs,
+    )
+
+
+def bgmv_expand_slice(
+    inputs: torch.Tensor,
+    lora_b_weights: torch.Tensor,
+    output_tensor: torch.Tensor,
+    lora_indices_tensor: torch.Tensor,
+    slice_offset: int,
+    slice_size: int,
+    add_inputs: bool = True,
+):
+    """
+    Args:
+        inputs (torch.Tensor): Shape: `[batch_size, hidden_size]`.
+        lora_b_weights (torch.Tensor): Shape: `[lora_num, rank, hidden_size]`.
+        output_tensor (torch.Tensor): Shape: `[batch_size, rank]`.
+        lora_indices_tensor (torch.Tensor): Shape: `[batch_size]`.
+            The LoRA index
+            corresponding to each batch. An index of -1 means no lora should be
+            applied.
+        slice_offset (int): output_tensor's offset
+        slice_size (int): current output_tensor's size
+        add_inputs (bool, optional):  Defaults to False. adds the final lora
+            results to the output.
+
+    Semantics:
+      for i in range(inputs.size(0)):
+        output_tensor[i][slice_offset:slice_offset+slice_size] =
+            inputs[i] @ lora_b_weights[lora_indices_tensor[i]]
+            + (inputs[i] if add_inputs else 0)
+    """
+    torch.ops._xpu_C.bgmv_expand_slice(
+        output_tensor,
+        inputs,
+        lora_b_weights,
+        lora_indices_tensor,
+        slice_offset,
+        add_inputs,
+    )

csrc/xpu/torch_bindings.cpp

@@ -5,7 +5,6 @@
 import torch
 import vllm_xpu_kernels._C  # noqa: F401
 import vllm_xpu_kernels._moe_C  # noqa: F401
-import vllm_xpu_kernels._lora_C  # noqa: F401
 
 
 # layer norm ops
@@ -162,91 +161,3 @@ def swigluoai_and_mul(
 # moe
 def moe_sum(input: torch.Tensor, output: torch.Tensor) -> None:
     torch.ops._moe_C.moe_sum(input, output)
-
-
-def bgmv_shrink(
-    inputs: torch.Tensor,
-    lora_a_weights: torch.Tensor,
-    output_tensor: torch.Tensor,
-    lora_indices_tensor: torch.Tensor,
-    scaling: float = 1.0,
-) -> None:
-    torch.ops._lora_C.bgmv_shrink(
-        output_tensor,
-        inputs,
-        lora_a_weights,
-        lora_indices_tensor,
-        scaling,
-    )
-
-
-def bgmv_expand(
-    inputs: torch.Tensor,
-    lora_b_weights: torch.Tensor,
-    output_tensor: torch.Tensor,
-    lora_indices_tensor: torch.Tensor,
-    add_inputs: bool = True,
-):
-    """
-    Args:
-        inputs (torch.Tensor): Shape: `[batch_size, hidden_size]`.
-        lora_b_weights (torch.Tensor): Shape: `[lora_num, rank, hidden_size]`.
-        output_tensor (torch.Tensor): Shape: `[batch_size, rank]`.
-        lora_indices_tensor (torch.Tensor): Shape: `[batch_size]`.
-         The LoRA index corresponding to each batch. An index of -1 means
-            no lora should be applied.
-        add_inputs (bool, optional):  Defaults to False. adds the final lora
-            results to the output.
-
-    Semantics:
-      for i in range(inputs.size(0)):
-        output_tensor[i] =
-            inputs[i] @ lora_b_weights[lora_indices_tensor[i]]
-            + (inputs[i] if add_inputs else 0)
-    """
-    torch.ops._lora_C.bgmv_expand(
-        output_tensor,
-        inputs,
-        lora_b_weights,
-        lora_indices_tensor,
-        add_inputs,
-    )
-
-
-def bgmv_expand_slice(
-    inputs: torch.Tensor,
-    lora_b_weights: torch.Tensor,
-    output_tensor: torch.Tensor,
-    lora_indices_tensor: torch.Tensor,
-    slice_offset: int,
-    slice_size: int,
-    add_inputs: bool = True,
-):
-    """
-    Args:
-        inputs (torch.Tensor): Shape: `[batch_size, hidden_size]`.
-        lora_b_weights (torch.Tensor): Shape: `[lora_num, rank, hidden_size]`.
-        output_tensor (torch.Tensor): Shape: `[batch_size, rank]`.
-        lora_indices_tensor (torch.Tensor): Shape: `[batch_size]`.
-            The LoRA index
-            corresponding to each batch. An index of -1 means no lora should be
-            applied.
-        slice_offset (int): output_tensor's offset
-        slice_size (int): current output_tensor's size
-        add_inputs (bool, optional):  Defaults to False. adds the final lora
-            results to the output.
-
-    Semantics:
-      for i in range(inputs.size(0)):
-        output_tensor[i][slice_offset:slice_offset+slice_size] =
-            inputs[i] @ lora_b_weights[lora_indices_tensor[i]]
-            + (inputs[i] if add_inputs else 0)
-    """
-    torch.ops._lora_C.bgmv_expand_slice(
-        output_tensor,
-        inputs,
-        lora_b_weights,
-        lora_indices_tensor,
-        slice_offset,
-        add_inputs,
-    )