refactor: using tvm-ffi for multi-platform bindings (#1641)

<!-- .github/pull_request_template.md -->

## 📌 Description

This PR refactors the codebase to use tvm-ffi for python bindings. The
goal is as following:
1. Supporting multiple backends
2. PyTorch version agnostic aot wheels
3. Accelerate compilation time and reduce binary bloat by minimizing
dependencies

Co-authored-by: Yaxing Cai<[email protected]>
Co-authored-by: Kathryn Chen<[email protected]>

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

cc @tqchen

---------

Co-authored-by: Yaxing Cai <[email protected]>
Co-authored-by: Kathryn-cat <[email protected]>
Co-authored-by: Kathryn-cat <[email protected]>
Co-authored-by: Yaxing Cai <[email protected]>

Author: 3 people

ci/bash.sh

@@ -70,6 +70,10 @@ echo "ENV VARIABLES: ${DOCKER_ENV}"
 echo "VOLUMES: ${DOCKER_VOLUMNS}"
 echo "COMMANDS: '${COMMAND[@]}'"
 
+# Pull the latest docker image
+echo "Pulling latest image: ${DOCKER_IMAGE_NAME}"
+${DOCKER_BINARY} pull ${DOCKER_IMAGE_NAME}
+
 # By default we cleanup - remove the container once it finish running (--rm)
 # and share the PID namespace (--pid=host) so the process inside does not have
 # pid 1 and SIGKILL is propagated to the process inside (jenkins can kill it).

csrc/activation.cu

@@ -17,14 +17,15 @@
 #include <flashinfer/attention/scheduler.cuh>
 #include <flashinfer/layout.cuh>
 #include <flashinfer/pos_enc.cuh>
-#include <optional>
 
 #include "batch_attention_config.inc"
-#include "pytorch_conversion_utils.h"
-#include "pytorch_extension_utils.h"
+#include "tvm_ffi_utils.h"
 
 namespace flashinfer {
 
+using tvm::ffi::Array;
+using tvm::ffi::Optional;
+
 template <uint32_t CTA_TILE_Q_1, uint32_t CTA_TILE_Q_2, uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_VO,
           MaskMode MASK_MODE, typename AttentionVariant, typename Params>
 cudaError_t BatchPagedAttentionPersistent(const Params params_1, const Params params_2,
@@ -34,80 +35,73 @@ cudaError_t BatchPagedAttentionPersistent(const Params params_1, const Params pa
 
 using namespace flashinfer;
 
-at::Tensor BatchPagedAttentionPlan(at::Tensor float_workspace_buffer,
-                                   at::Tensor int_workspace_buffer,
-                                   at::Tensor page_locked_int_workspace_buffer,
-                                   at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor kv_len,
-                                   int64_t batch_size, int64_t num_qo_heads, int64_t num_kv_heads,
-                                   int64_t head_dim_o, bool causal) {
+Array<int64_t> BatchPagedAttentionPlan(Tensor float_workspace_buffer, Tensor int_workspace_buffer,
+                                       Tensor page_locked_int_workspace_buffer, Tensor qo_indptr,
+                                       Tensor kv_indptr, Tensor kv_len, int64_t batch_size,
+                                       int64_t num_qo_heads, int64_t num_kv_heads,
+                                       int64_t head_dim_o, bool causal) {
   size_t float_workspace_size_in_bytes =
-      float_workspace_buffer.size(0) * float_workspace_buffer.element_size();
+      float_workspace_buffer->shape[0] * get_element_size(float_workspace_buffer);
   size_t int_workspace_size_in_bytes =
-      int_workspace_buffer.size(0) * int_workspace_buffer.element_size();
+      int_workspace_buffer->shape[0] * get_element_size(int_workspace_buffer);
 
   HolisticPlanInfo<2> plan_info;
 
-  const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
-  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  cudaSetDevice(float_workspace_buffer->device.device_id);
+  const cudaStream_t stream = get_stream(float_workspace_buffer->device);
 
   cudaError_t status = TwoStageHolisticPlan<IdType>(
-      float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes,
-      int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(),
-      int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr<IdType>(),
-      kv_indptr.data_ptr<IdType>(), kv_len.data_ptr<IdType>(), batch_size, num_qo_heads,
-      num_kv_heads, head_dim_o, causal, stream);
+      float_workspace_buffer->data, float_workspace_size_in_bytes, int_workspace_buffer->data,
+      page_locked_int_workspace_buffer->data, int_workspace_size_in_bytes, plan_info,
+      static_cast<IdType*>(qo_indptr->data), static_cast<IdType*>(kv_indptr->data),
+      static_cast<IdType*>(kv_len->data), batch_size, num_qo_heads, num_kv_heads, head_dim_o,
+      causal, stream);
 
-  TORCH_CHECK(status == cudaSuccess,
-              "Failed to plan persistent paged attention, error: ", cudaGetErrorString(status));
+  TVM_FFI_ICHECK(status == cudaSuccess)
+      << "Failed to plan persistent paged attention, error: " << cudaGetErrorString(status);
 
-  return vec_to_tensor(plan_info.ToVector());
+  return Array(plan_info.ToVector());
 }
 
-void BatchPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
-                            at::Tensor plan_info_vec, at::Tensor q, at::Tensor k_cache,
-                            at::Tensor v_cache, at::Tensor kv_indices, at::Tensor o,
-                            std::optional<at::Tensor> maybe_lse, int64_t mask_mode_code,
-                            int64_t layout_code, int64_t num_qo_heads, int64_t num_kv_heads,
-                            int64_t page_size,
+void BatchPagedAttentionRun(Tensor float_workspace_buffer, Tensor int_workspace_buffer,
+                            Array<int64_t> plan_info_vec, Tensor q, Tensor k_cache, Tensor v_cache,
+                            Tensor kv_indices, Tensor o, Optional<Tensor> maybe_lse,
+                            int64_t mask_mode_code, int64_t layout_code, int64_t num_qo_heads,
+                            int64_t num_kv_heads, int64_t page_size,
                             double v_scale,  // must use double due to pytorch binding
                             double sm_scale,
                             double logits_soft_cap ADDITIONAL_FUNC_PARAMS PROFILER_FUNC_PARAMS) {
   HolisticPlanInfo<2> plan_info;
-  plan_info.FromVector(tensor_to_vec(plan_info_vec));
-
-  auto device = q.device();
+  plan_info.FromVector(std::vector<int64_t>(plan_info_vec.begin(), plan_info_vec.end()));
 
-  void* float_buffer_ptr = float_workspace_buffer.data_ptr();
-  void* int_buffer_ptr = int_workspace_buffer.data_ptr();
+  void* float_buffer_ptr = float_workspace_buffer->data;
+  void* int_buffer_ptr = int_workspace_buffer->data;
 
   const MaskMode mask_mode = static_cast<MaskMode>(mask_mode_code);
 
-  auto q_scalar_type = q.scalar_type();
-  auto kv_scalar_type = k_cache.scalar_type();
-
   // NOTE (Yilong): assume both q and o are NHD
-  unsigned int q_stride_n = q.stride(0);
-  unsigned int q_stride_h = q.stride(1);
+  unsigned int q_stride_n = q->strides[0];
+  unsigned int q_stride_h = q->strides[1];
 
   // layout only constraint paged KV
   const QKVLayout kv_layout = static_cast<QKVLayout>(layout_code);
-  unsigned int k_stride_page = k_cache.stride(0);
-  unsigned int v_stride_page = v_cache.stride(0);
+  unsigned int k_stride_page = k_cache->strides[0];
+  unsigned int v_stride_page = v_cache->strides[0];
   unsigned int k_stride_n, k_stride_h, v_stride_n, v_stride_h;
   if (kv_layout == QKVLayout::kNHD) {
-    k_stride_h = k_cache.stride(2);
-    k_stride_n = k_cache.stride(1);
-    v_stride_h = v_cache.stride(2);
-    v_stride_n = v_cache.stride(1);
+    k_stride_h = k_cache->strides[2];
+    k_stride_n = k_cache->strides[1];
+    v_stride_h = v_cache->strides[2];
+    v_stride_n = v_cache->strides[1];
   } else {
-    k_stride_h = k_cache.stride(1);
-    k_stride_n = k_cache.stride(2);
-    v_stride_h = v_cache.stride(1);
-    v_stride_n = v_cache.stride(2);
+    k_stride_h = k_cache->strides[1];
+    k_stride_n = k_cache->strides[2];
+    v_stride_h = v_cache->strides[1];
+    v_stride_n = v_cache->strides[2];
   }
 
-  const c10::cuda::OptionalCUDAGuard device_guard(device);
-  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  cudaSetDevice(q->device.device_id);
+  const cudaStream_t stream = get_stream(q->device);
 
   DISPATCH_context(
       DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
@@ -116,17 +110,17 @@ void BatchPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int_wo
         IdType* len_kv_chunk =
             GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.len_kv_chunk_offset);
         for (int i = 0; i < 2; i++) {
-          params[i].q = static_cast<DTypeQ*>(q.data_ptr());
-          params[i].k = static_cast<DTypeKV*>(k_cache.data_ptr());
-          params[i].v = static_cast<DTypeKV*>(v_cache.data_ptr());
+          params[i].q = static_cast<DTypeQ*>(q->data);
+          params[i].k = static_cast<DTypeKV*>(k_cache->data);
+          params[i].v = static_cast<DTypeKV*>(v_cache->data);
 
           params[i].q_indptr =
               GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.tasks[i].q_indptr_offset);
           params[i].kv_indptr =
               GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.tasks[i].kv_indptr_offset);
           params[i].partial_indptr = GetPtrFromBaseOffset<IdType>(
               int_buffer_ptr, plan_info.tasks[i].partial_indptr_offset);
-          params[i].kv_indices = static_cast<int*>(kv_indices.data_ptr());
+          params[i].kv_indices = static_cast<int*>(kv_indices->data);
           params[i].q_len =
               GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.tasks[i].q_len_offset);
           params[i].kv_len =
@@ -143,9 +137,9 @@ void BatchPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int_wo
               GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.tasks[i].work_indptr_offset);
           params[i].len_kv_chunk = len_kv_chunk + i;
 
-          params[i].final_o = static_cast<DTypeO*>(o.data_ptr());
+          params[i].final_o = static_cast<DTypeO*>(o->data);
           params[i].final_lse =
-              maybe_lse.has_value() ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+              maybe_lse.has_value() ? static_cast<float*>(maybe_lse.value()->data) : nullptr;
           params[i].partial_o =
               GetPtrFromBaseOffset<DTypeO>(float_buffer_ptr, plan_info.partial_o_offset);
           params[i].partial_lse =
@@ -184,8 +178,8 @@ void BatchPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int_wo
         cudaError_t status = BatchPagedAttentionPersistent<128, 16, HEAD_DIM_QK, HEAD_DIM_VO,
                                                            MASK_MODE, AttentionVariant>(
             params[0], params[1], plan_info.num_blks_x, plan_info.num_blks_y, stream);
-        TORCH_CHECK(status == cudaSuccess, "Failed to run persistent paged attention, error: ",
-                    cudaGetErrorString(status));
+        TVM_FFI_ICHECK(status == cudaSuccess)
+            << "Failed to run persistent paged attention, error: " << cudaGetErrorString(status);
         return true;
       });
 }