From 9b432e16ba45e556f1fd8d9be4766ddf972f1768 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Fri, 6 Jun 2025 00:00:09 -0400 Subject: [PATCH] [Backend] JIT compile FlashInfer kernel with FFI header This PR adds the FFI header directory for FlashInfer JIT compilation. --- python/tvm/relax/backend/cuda/flashinfer.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/tvm/relax/backend/cuda/flashinfer.py b/python/tvm/relax/backend/cuda/flashinfer.py index 687987e4d66b..47a4946ca97d 100644 --- a/python/tvm/relax/backend/cuda/flashinfer.py +++ b/python/tvm/relax/backend/cuda/flashinfer.py @@ -16,13 +16,13 @@ # under the License. """FlashInfer JIT compilation module for CUDA backend""" +import hashlib +import json import os import subprocess from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import List -import hashlib -import json import tvm from tvm.target import Target @@ -130,6 +130,7 @@ def get_object_file_path(src: Path) -> Path: FLASHINFER_CSRC_DIR, FLASHINFER_TVM_BINDING_DIR, Path(tvm_home).resolve() / "include", + Path(tvm_home).resolve() / "ffi" / "include", Path(tvm_home).resolve() / "3rdparty" / "dlpack" / "include", Path(tvm_home).resolve() / "3rdparty" / "dmlc-core" / "include", ] + CUTLASS_INCLUDE_DIRS