Skip to content

refactor: cleanup codebase after tvm-ffi refactor#1795

Merged
yzh119 merged 18 commits intoflashinfer-ai:mainfrom
yzh119:ffi-bump-followup
Sep 29, 2025
Merged

refactor: cleanup codebase after tvm-ffi refactor#1795
yzh119 merged 18 commits intoflashinfer-ai:mainfrom
yzh119:ffi-bump-followup

Conversation

@yzh119
Copy link
Copy Markdown
Collaborator

@yzh119 yzh119 commented Sep 27, 2025

📌 Description

The codegen logic for pytorch and tvm should unify after #1641 , and this PR cleans up the related codegen functions in tvm_bindings.

Other changes:

  1. update tvm-ffi to 0.1.0b11 to incorporate [ABI] Introduce ShapeView Minimize TensorObj exposure apache/tvm-ffi#67 and [CYTHON] Fix stream passing bug apache/tvm-ffi#68
  2. rename of source files: _ops.cu and _pybind.cu renamed to _binding.cu
  3. remove torch related header include/library linking in ninja files (Support PaddlePaddle with compatible API #1642 (comment))
  4. remove the use of use_torch_stream in unittests, they are no longer required after [CYTHON] Fix stream passing bug apache/tvm-ffi#68

🔍 Related Issues

#1641

🚀 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

  • I have installed pre-commit by running pip install pre-commit (or used your preferred method).
  • I have installed the hooks with pre-commit install.
  • 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.

🧪 Tests

  • Tests have been added or updated as needed.
  • All tests are passing (unittest, etc.).

Reviewer Notes

cc @MasterJH5574 please let us know what changes do we need to make to help you bump to the latest version of flashinfer in MLC.

@gemini-code-assist
Copy link
Copy Markdown
Contributor

Summary of Changes

Hello @yzh119, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request primarily focuses on refactoring and cleaning up the codebase following a previous tvm-ffi refactor. The core intent is to unify the codegen logic between PyTorch and TVM, leading to a more streamlined and maintainable project structure. This involves updating the tvm-ffi dependency, standardizing naming conventions for C++ binding files, and removing obsolete TVM-specific binding implementations. Additionally, the build system has been simplified by removing redundant Torch-related configurations.

Highlights

  • Codebase Unification: The codegen logic for PyTorch and TVM has been unified, leading to a cleaner and more consistent codebase.
  • Dependency Update: The apache-tvm-ffi dependency has been updated to version 0.1.0b10 to incorporate recent upstream changes.
  • File Renaming Convention: Numerous C++ source files (.cu) have been renamed from _ops.cu or _pybind.cu to a standardized _binding.cu suffix, reflecting their role in providing bindings.
  • Removal of TVM Bindings: All TVM-specific binding files and related Python code have been removed, streamlining the project structure.
  • Build System Cleanup: The build system configuration in flashinfer/jit/cpp_ext.py has been simplified by removing Pybind11 ABI flags and Torch-specific include/linker paths, and updating the nvcc environment variable.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request is a significant refactoring to unify the codegen logic for PyTorch and TVM by moving to tvm-ffi. This involves renaming many files from _ops.cu and _pybind.cu to _binding.cu, and removing the old tvm_binding directory. The changes in flashinfer/jit/cpp_ext.py correctly remove dependencies on PyTorch's C++ extension build system, which is a good step towards a more backend-agnostic FFI approach. The overall refactoring seems well-executed. However, I've found a systematic typo in the Python files that reference the newly renamed CUDA files: they use _bindings.cu (plural) while the renamed files are _binding.cu (singular). This will cause build failures and needs to be corrected across all affected files. I've added comments for each occurrence to help you fix them.

@yzh119 yzh119 enabled auto-merge (squash) September 28, 2025 19:08
@yzh119 yzh119 merged commit d161b56 into flashinfer-ai:main Sep 29, 2025
3 checks passed
@MasterJH5574
Copy link
Copy Markdown
Collaborator

All necessary changes on flashinfer side to support tvm and mlc are included in #1836.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants