hotfix: fix rope tvm wrapper (#601)

The TVM wrapper was broken in #599 because of API changes, this PR fixes the issue.
flashinfer-ai · Nov 10, 2024 · 3dd9405 · 3dd9405
1 parent eb9bc71
commit 3dd9405
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/tvm_wrapper.cu b/src/tvm_wrapper.cu
@@ -692,7 +692,8 @@ void _FlashInferBatchQKApplyRotaryInPlace(DLTensor* q, DLTensor* k, DLTensor* in
         cudaError_t status = BatchQKApplyRotaryInPlace(
             static_cast<dtype*>(q->data), static_cast<dtype*>(k->data),
             static_cast<idtype*>(indptr->data), static_cast<idtype*>(offsets->data), batch_size,
-            num_qo_heads, num_kv_heads, head_dim, q_stride_n, q_stride_h, k_stride_n, k_stride_h,
+            num_qo_heads, num_kv_heads, /*rotary_dim=*/head_dim, head_dim, q_stride_n, q_stride_h,
+            k_stride_n, k_stride_h,
             /*interleave=*/false, rope_scale, rope_theta);
         if (status != cudaSuccess) {
           LOG(FATAL) << "FlashInfer CUDA kernel error " << cudaGetErrorString(status);