refactor: stop build flash_infer kernel (#386)

Author: guocuimi

scalellm/CMakeLists.txt

@@ -18,7 +18,6 @@ pybind_extension(
   DEPS
     :llm_handler
     :marlin.kernels
-    :flash_infer.kernels
     torch
     torch_python
     absl::strings

scalellm/csrc/kernels.cu

@@ -1,7 +1,6 @@
 #include <pybind11/pybind11.h>
 #include <torch/extension.h>
 
-#include "kernels/attention/flash_infer/attention_wrapper.h"
 #include "kernels/quantization/marlin.h"
 
 namespace llm::csrc {
@@ -53,16 +52,6 @@ void init_kernels(py::module_& m) {
         py::arg("q_weight"),
         py::arg("out"),
         py::arg("num_bits"));
-
-  // flashinfer kernels
-  py::class_<flashinfer::BatchPrefillWrapper>(m, "BatchPrefillWrapper")
-      .def(py::init<bool>())
-      .def("plan", &flashinfer::BatchPrefillWrapper::Plan)
-      .def("is_cuda_graph_enabled",
-           &flashinfer::BatchPrefillWrapper::IsCUDAGraphEnabled)
-      .def("update_page_locked_buffer_size",
-           &flashinfer::BatchPrefillWrapper::UpdatePageLockedBufferSize)
-      .def("run", &flashinfer::BatchPrefillWrapper::Run);
 }
 
 }  // namespace llm::csrc

src/kernels/attention/CMakeLists.txt

@@ -80,5 +80,5 @@ cc_binary(
 )
 
 add_subdirectory(flash_attn)
-add_subdirectory(flash_infer)
+# add_subdirectory(flash_infer)
 add_subdirectory(tools)

tests/kernels/attention/flash_infer_kv_fp8_test.py

@@ -7,6 +7,7 @@
 import scalellm._C.kernels as kernels  # type: ignore
 
 
+@pytest.mark.skip(reason="Not implemented")
 @pytest.mark.parametrize("seq_lens", [[(1, 100)], [(100, 100)], [(1, 100), (15, 15), (111, 234), (1000, 10000)]])
 @pytest.mark.parametrize("num_heads", [(8, 8), (8, 4), (8, 2), (8, 1)])
 @pytest.mark.parametrize("head_size", [64, 128, 256])

tests/kernels/attention/flash_infer_test.py

@@ -7,6 +7,7 @@
 import scalellm._C.kernels as kernels  # type: ignore
 
 
+@pytest.mark.skip(reason="Not implemented")
 @pytest.mark.parametrize("seq_lens", [[(1, 100)], [(100, 100)], [(1, 100), (15, 15), (111, 234), (1000, 10000)]])
 @pytest.mark.parametrize("num_heads", [(8, 8), (8, 4), (8, 2), (8, 1)])
 @pytest.mark.parametrize("head_size", [64, 128, 256])