[JIT] Fix MLA header in TVM binding

This PR fixes the header include, following changes in #887.
flashinfer-ai · Feb 23, 2025 · 649475a · 649475a
1 parent 2b24293
commit 649475a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tvm_binding/batch_mla_run.cu b/tvm_binding/batch_mla_run.cu
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#include <flashinfer/attention/mla_fa2.cuh>
+#include <flashinfer/attention/mla.cuh>
 #include <flashinfer/attention/scheduler.cuh>
 #include <flashinfer/fastdiv.cuh>
 #include <optional>